package vidmod import ( "context" "errors" "fmt" "reflect" "sort" "strings" "sync" "sync/atomic" "testing" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" ) type fakeRecommendScoreStore struct { maxEligibleID func(context.Context) (ObjectID, error) countUninitialized func(context.Context, ObjectID, int64) (int64, error) openCursor func(context.Context, ObjectID, int32) (recommendScoreCursor, error) findBatch func(context.Context, ObjectID, ObjectID, int) ([]recommendScoreDocument, error) bulkWrite func(context.Context, []mongo.WriteModel) error } func (f fakeRecommendScoreStore) MaxEligibleID(ctx context.Context) (ObjectID, error) { return f.maxEligibleID(ctx) } func (f fakeRecommendScoreStore) CountUninitialized( ctx context.Context, maxID ObjectID, limit int64, ) (int64, error) { if f.countUninitialized == nil { return 0, nil } return f.countUninitialized(ctx, maxID, limit) } func (f fakeRecommendScoreStore) OpenCursor( ctx context.Context, maxID ObjectID, batchSize int32, ) (recommendScoreCursor, error) { if f.openCursor != nil { return f.openCursor(ctx, maxID, batchSize) } if f.findBatch == nil { return nil, errors.New("fake recommend score cursor is not configured") } documents := make([]recommendScoreDocument, 0) afterID := ObjectID{} limit := int(batchSize) for { batch, err := f.findBatch(ctx, afterID, maxID, limit) if err != nil { return nil, err } if len(batch) == 0 { break } documents = append(documents, batch...) afterID = batch[len(batch)-1].ID if len(batch) < limit || afterID == maxID { break } } return &fakeRecommendScoreCursor{documents: documents, current: -1}, nil } func (f fakeRecommendScoreStore) BulkWrite(ctx context.Context, writes []mongo.WriteModel) error { return f.bulkWrite(ctx, writes) } type fakeRecommendScoreCursor struct { documents []recommendScoreDocument index int current int cursorErr error decodeErrAt int decodeErr error closeErr error closeCalls atomic.Int32 closeContextCanceled atomic.Bool } func (f *fakeRecommendScoreCursor) Next(ctx context.Context) bool { if err := ctx.Err(); err != nil { f.cursorErr = err return false } if f.index >= len(f.documents) { return false } f.current = f.index f.index++ return true } func (f *fakeRecommendScoreCursor) Decode(value interface{}) error { if f.decodeErr != nil && f.current == f.decodeErrAt { return f.decodeErr } document, ok := value.(*recommendScoreDocument) if !ok { return fmt.Errorf("unexpected decode target %T", value) } if f.current < 0 || f.current >= len(f.documents) { return errors.New("decode called without current document") } *document = f.documents[f.current] return nil } func (f *fakeRecommendScoreCursor) Err() error { return f.cursorErr } func (f *fakeRecommendScoreCursor) Close(ctx context.Context) error { f.closeCalls.Add(1) f.closeContextCanceled.Store(ctx.Err() != nil) return f.closeErr } type refreshRecommendScoresResult struct { candidates []RecommendCandidate err error } func recommendBatchTestID(value byte) ObjectID { var id ObjectID id[len(id)-1] = value return id } func recommendBatchTestScore(value int64) *int64 { return &value } func recommendBatchTestDocuments(count int) []recommendScoreDocument { documents := make([]recommendScoreDocument, count) for i := range documents { documents[i] = recommendScoreDocument{ ID: recommendBatchTestID(byte(i + 1)), RecommendLikeCount: int64(i + 1), RecommendInitialized: true, } } return documents } func recommendBatchTestFind( documents []recommendScoreDocument, afterID, maxID ObjectID, limit int, ) []recommendScoreDocument { batch := make([]recommendScoreDocument, 0, limit) for _, document := range documents { if compareObjectID(document.ID, afterID) <= 0 || compareObjectID(document.ID, maxID) > 0 { continue } batch = append(batch, document) if len(batch) == limit { break } } return batch } func TestRefreshRecommendScoresBatchesAtFixedHighWater(t *testing.T) { const ( batchSize = 2 cursorBatchSize = 4 ) initialDocuments := recommendBatchTestDocuments(5) highWater := initialDocuments[len(initialDocuments)-1].ID documents := append([]recommendScoreDocument(nil), initialDocuments...) var stateMu sync.Mutex maxEligibleIDCalls := 0 openCursorCalls := 0 var openedMaxID ObjectID var openedBatchSize int32 var cursor *fakeRecommendScoreCursor bulkSizes := make([]int, 0, 3) store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { stateMu.Lock() defer stateMu.Unlock() maxEligibleIDCalls++ return highWater, nil }, openCursor: func( _ context.Context, maxID ObjectID, serverBatchSize int32, ) (recommendScoreCursor, error) { stateMu.Lock() defer stateMu.Unlock() openCursorCalls++ openedMaxID = maxID openedBatchSize = serverBatchSize // 模拟扫描开始后插入的数据;固定高水位必须把它排除在本轮之外。 documents = append(documents, recommendScoreDocument{ ID: recommendBatchTestID(6), RecommendLikeCount: 6, RecommendInitialized: true, }) filtered := make([]recommendScoreDocument, 0, len(documents)) for _, document := range documents { if compareObjectID(document.ID, maxID) <= 0 { filtered = append(filtered, document) } } cursor = &fakeRecommendScoreCursor{documents: filtered, current: -1} return cursor, nil }, bulkWrite: func(_ context.Context, writes []mongo.WriteModel) error { stateMu.Lock() defer stateMu.Unlock() bulkSizes = append(bulkSizes, len(writes)) return nil }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), batchSize, cursorBatchSize, 2, store, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } stateMu.Lock() gotMaxEligibleIDCalls := maxEligibleIDCalls gotOpenCursorCalls := openCursorCalls gotOpenedMaxID := openedMaxID gotOpenedBatchSize := openedBatchSize gotBulkSizes := append([]int(nil), bulkSizes...) stateMu.Unlock() if gotMaxEligibleIDCalls != 1 { t.Fatalf("MaxEligibleID calls = %d, want 1", gotMaxEligibleIDCalls) } if gotOpenCursorCalls != 1 { t.Fatalf("OpenCursor calls = %d, want 1", gotOpenCursorCalls) } if gotOpenedMaxID != highWater { t.Fatalf("OpenCursor maxID = %s, want fixed high water %s", gotOpenedMaxID.Hex(), highWater.Hex()) } if gotOpenedBatchSize != cursorBatchSize { t.Fatalf("OpenCursor batch size = %d, want %d", gotOpenedBatchSize, cursorBatchSize) } if cursor == nil { t.Fatal("OpenCursor returned a nil cursor") } if cursor.closeCalls.Load() != 1 { t.Fatalf("Cursor close calls = %d, want 1", cursor.closeCalls.Load()) } sort.Ints(gotBulkSizes) if got, want := gotBulkSizes, []int{1, 2, 2}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] { t.Fatalf("BulkWrite batch sizes = %v, want %v", got, want) } if len(candidates) != len(initialDocuments) { t.Fatalf("candidate count = %d, want %d", len(candidates), len(initialDocuments)) } seen := make(map[ObjectID]bool, len(candidates)) for _, candidate := range candidates { seen[candidate.ID] = true } for _, document := range initialDocuments { if !seen[document.ID] { t.Errorf("candidate %s is missing", document.ID.Hex()) } } if seen[recommendBatchTestID(6)] { t.Error("candidate inserted above the fixed high water was included") } } func TestPrepareRecommendScoreBatchWritesOnlyChangedOrUninitialized(t *testing.T) { now := time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC) documents := []recommendScoreDocument{ { ID: recommendBatchTestID(1), LikeCount: 1, RecommendLikeCount: 10, RecommendCollectCount: 2, RecommendCommentCount: 3, RecommendShareCount: 4, RecommendScore: recommendBatchTestScore(CalculateRecommendScore(10, 2, 3, 4)), RecommendInitialized: true, }, { ID: recommendBatchTestID(2), LikeCount: 8, RecommendLikeCount: 5, RecommendCollectCount: 1, RecommendScore: recommendBatchTestScore(0), RecommendInitialized: true, }, { ID: recommendBatchTestID(3), LikeCount: 7, RecommendScore: recommendBatchTestScore(7), RecommendInitialized: false, }, { ID: recommendBatchTestID(4), RecommendScore: recommendBatchTestScore(0), RecommendInitialized: true, }, { ID: recommendBatchTestID(5), RecommendScore: nil, RecommendInitialized: true, }, } candidates, writes := prepareRecommendScoreBatch(documents, now) if len(candidates) != len(documents) { t.Fatalf("candidate count = %d, want %d", len(candidates), len(documents)) } if len(writes) != 3 { t.Fatalf("write count = %d, want 3 changed/uninitialized/missing-score documents", len(writes)) } wantWriteIDs := []ObjectID{documents[1].ID, documents[2].ID, documents[4].ID} for i, write := range writes { update, ok := write.(*mongo.UpdateOneModel) if !ok { t.Fatalf("write %d type = %T, want *mongo.UpdateOneModel", i, write) } filter, ok := update.Filter.(bson.M) if !ok || filter["_id"] != wantWriteIDs[i] { t.Fatalf("write %d filter = %#v, want _id %s", i, update.Filter, wantWriteIDs[i].Hex()) } pipeline, ok := update.Update.(mongo.Pipeline) if !ok { t.Fatalf("write %d update type = %T, want mongo.Pipeline", i, update.Update) } pipelineText := fmt.Sprint(pipeline) for _, field := range []string{ "$recommendLikeCount", "$recommendCollectCount", "$recommendCommentCount", "$recommendShareCount", } { if !strings.Contains(pipelineText, field) { t.Fatalf("write %d pipeline does not reference %s: %s", i, field, pipelineText) } } } wantScores := []int64{ CalculateRecommendScore(10, 2, 3, 4), CalculateRecommendScore(8, 1, 0, 0), 7, 0, 0, } for i, candidate := range candidates { if candidate.RecommendScore != wantScores[i] { t.Errorf("candidate %d score = %d, want %d", i, candidate.RecommendScore, wantScores[i]) } } } func TestRefreshRecommendScoresSkipsBulkWriteForUnchangedBatches(t *testing.T) { documents := recommendBatchTestDocuments(5) for i := range documents { documents[i].RecommendScore = recommendBatchTestScore(CalculateRecommendScore( documents[i].RecommendLikeCount, documents[i].RecommendCollectCount, documents[i].RecommendCommentCount, documents[i].RecommendShareCount, )) } highWater := documents[len(documents)-1].ID var bulkCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { bulkCalls.Add(1) return errors.New("unchanged batch must not be written") }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 2, 2, store, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } if len(candidates) != len(documents) { t.Fatalf("candidate count = %d, want %d", len(candidates), len(documents)) } if got := bulkCalls.Load(); got != 0 { t.Fatalf("BulkWrite calls = %d, want 0 for unchanged documents", got) } } func TestRefreshRecommendScoresInitializationGuardStopsBeforeWrites(t *testing.T) { documents := recommendBatchTestDocuments(2) highWater := documents[len(documents)-1].ID var cursorCalls, bulkCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, countUninitialized: func( _ context.Context, gotMaxID ObjectID, limit int64, ) (int64, error) { if gotMaxID != highWater { t.Fatalf("CountUninitialized maxID = %s, want %s", gotMaxID.Hex(), highWater.Hex()) } if limit != 11 { t.Fatalf("CountUninitialized limit = %d, want 11", limit) } return 11, nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { cursorCalls.Add(1) return nil, errors.New("cursor must not open after guard rejection") }, bulkWrite: func(context.Context, []mongo.WriteModel) error { bulkCalls.Add(1) return errors.New("bulk write must not run after guard rejection") }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 2, 2, store, 10, ) var limitErr RecommendInitializationLimitError if !errors.As(err, &limitErr) { t.Fatalf("error = %v, want RecommendInitializationLimitError", err) } if limitErr.Limit != 10 || limitErr.ObservedAtLeast != 11 { t.Fatalf("limit error = %+v", limitErr) } if candidates != nil { t.Fatalf("candidates = %#v, want nil", candidates) } if cursorCalls.Load() != 0 || bulkCalls.Load() != 0 { t.Fatalf("cursorCalls=%d bulkCalls=%d, want zero", cursorCalls.Load(), bulkCalls.Load()) } } func TestRefreshRecommendScoresInitializationGuardAllowsExactLimit(t *testing.T) { documents := recommendBatchTestDocuments(2) for i := range documents { documents[i].RecommendScore = recommendBatchTestScore(documents[i].RecommendLikeCount) } highWater := documents[len(documents)-1].ID var countCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, countUninitialized: func(context.Context, ObjectID, int64) (int64, error) { countCalls.Add(1) return 10, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return errors.New("unchanged documents must not be written") }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 2, 2, store, 10, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } if countCalls.Load() != 1 || len(candidates) != len(documents) { t.Fatalf("countCalls=%d candidateCount=%d", countCalls.Load(), len(candidates)) } } func TestRefreshRecommendScoresWritesOnlyChangedMixedBatch(t *testing.T) { documents := recommendBatchTestDocuments(4) for i := range documents { documents[i].RecommendScore = recommendBatchTestScore(documents[i].RecommendLikeCount) } documents[2].RecommendScore = recommendBatchTestScore(0) highWater := documents[len(documents)-1].ID var stateMu sync.Mutex var bulkSizes []int store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(_ context.Context, writes []mongo.WriteModel) error { stateMu.Lock() defer stateMu.Unlock() bulkSizes = append(bulkSizes, len(writes)) return nil }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 2, 2, store, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } if len(candidates) != len(documents) { t.Fatalf("candidate count = %d, want %d", len(candidates), len(documents)) } stateMu.Lock() gotBulkSizes := append([]int(nil), bulkSizes...) stateMu.Unlock() if len(gotBulkSizes) != 1 || gotBulkSizes[0] != 1 { t.Fatalf("BulkWrite batch sizes = %v, want [1]", gotBulkSizes) } } func TestRefreshRecommendScoresBoundsAndUsesConcurrency(t *testing.T) { const workerCount = 3 documents := recommendBatchTestDocuments(8) highWater := documents[len(documents)-1].ID release := make(chan struct{}) started := make(chan struct{}, len(documents)) var stateMu sync.Mutex active := 0 maxActive := 0 bulkCalls := 0 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(ctx context.Context, _ []mongo.WriteModel) error { stateMu.Lock() active++ bulkCalls++ if active > maxActive { maxActive = active } stateMu.Unlock() defer func() { stateMu.Lock() active-- stateMu.Unlock() }() started <- struct{}{} select { case <-release: return nil case <-ctx.Done(): return ctx.Err() } }, } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() result := make(chan refreshRecommendScoresResult, 1) go func() { candidates, err := refreshRecommendScores( ctx, time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 1, 1, workerCount, store, ) result <- refreshRecommendScoresResult{candidates: candidates, err: err} }() for i := 0; i < workerCount; i++ { select { case <-started: case <-ctx.Done(): t.Fatalf("only %d/%d workers reached the barrier: %v", i, workerCount, ctx.Err()) } } stateMu.Lock() gotMaxActiveAtBarrier := maxActive stateMu.Unlock() if gotMaxActiveAtBarrier != workerCount { t.Fatalf("active BulkWrite calls at barrier = %d, want %d", gotMaxActiveAtBarrier, workerCount) } close(release) var got refreshRecommendScoresResult select { case got = <-result: case <-ctx.Done(): t.Fatalf("refreshRecommendScores() did not finish: %v", ctx.Err()) } if got.err != nil { t.Fatalf("refreshRecommendScores() error = %v", got.err) } if len(got.candidates) != len(documents) { t.Fatalf("candidate count = %d, want %d", len(got.candidates), len(documents)) } stateMu.Lock() gotActive := active gotMaxActive := maxActive gotBulkCalls := bulkCalls stateMu.Unlock() if gotActive != 0 { t.Errorf("active BulkWrite calls after return = %d, want 0", gotActive) } if gotMaxActive != workerCount { t.Errorf("maximum concurrent BulkWrite calls = %d, want %d", gotMaxActive, workerCount) } if gotBulkCalls != len(documents) { t.Errorf("BulkWrite calls = %d, want %d", gotBulkCalls, len(documents)) } } func TestRefreshRecommendScoresFirstErrorCancelsWorkers(t *testing.T) { const workerCount = 2 writeErr := errors.New("bulk write failed") documents := recommendBatchTestDocuments(8) highWater := documents[len(documents)-1].ID started := make(chan int32, workerCount) releaseFailure := make(chan struct{}) cancelObserved := make(chan struct{}) var cancelObservedOnce sync.Once var releaseOnce sync.Once release := func() { releaseOnce.Do(func() { close(releaseFailure) }) } defer release() var bulkCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(ctx context.Context, _ []mongo.WriteModel) error { call := bulkCalls.Add(1) started <- call if call == 1 { select { case <-releaseFailure: return writeErr case <-ctx.Done(): return ctx.Err() } } <-ctx.Done() cancelObservedOnce.Do(func() { close(cancelObserved) }) return ctx.Err() }, } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() result := make(chan refreshRecommendScoresResult, 1) go func() { candidates, err := refreshRecommendScores( ctx, time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 1, 1, workerCount, store, ) result <- refreshRecommendScoresResult{candidates: candidates, err: err} }() for i := 0; i < workerCount; i++ { select { case <-started: case <-ctx.Done(): t.Fatalf("only %d/%d writes reached the barrier: %v", i, workerCount, ctx.Err()) } } release() var got refreshRecommendScoresResult select { case got = <-result: case <-ctx.Done(): t.Fatalf("refreshRecommendScores() did not finish after first error: %v", ctx.Err()) } if !errors.Is(got.err, writeErr) { t.Fatalf("refreshRecommendScores() error = %v, want %v", got.err, writeErr) } if got.candidates != nil { t.Fatalf("candidates = %v, want nil on write error", got.candidates) } select { case <-cancelObserved: default: t.Fatal("peer worker did not observe cancellation from the first error") } if gotCalls := bulkCalls.Load(); gotCalls != workerCount { t.Fatalf("BulkWrite calls = %d, want %d blocked workers only", gotCalls, workerCount) } } func TestRefreshRecommendScoresEmptyHighWaterSkipsWork(t *testing.T) { var findCalls atomic.Int32 var bulkCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return ObjectID{}, nil }, findBatch: func( context.Context, ObjectID, ObjectID, int, ) ([]recommendScoreDocument, error) { findCalls.Add(1) return nil, nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { bulkCalls.Add(1) return nil }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 2, 2, store, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } if candidates == nil || len(candidates) != 0 { t.Fatalf("candidates = %#v, want non-nil empty slice", candidates) } if got := findCalls.Load(); got != 0 { t.Errorf("FindBatch calls = %d, want 0", got) } if got := bulkCalls.Load(); got != 0 { t.Errorf("BulkWrite calls = %d, want 0", got) } } func TestRefreshRecommendScoresPropagatesContextCancellation(t *testing.T) { documents := recommendBatchTestDocuments(2) highWater := documents[len(documents)-1].ID writeStarted := make(chan struct{}, 1) cancelObserved := make(chan struct{}) store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, findBatch: func( _ context.Context, afterID, maxID ObjectID, limit int, ) ([]recommendScoreDocument, error) { return recommendBatchTestFind(documents, afterID, maxID, limit), nil }, bulkWrite: func(ctx context.Context, _ []mongo.WriteModel) error { writeStarted <- struct{}{} <-ctx.Done() close(cancelObserved) return ctx.Err() }, } ctx, cancel := context.WithCancel(context.Background()) defer cancel() guard, stopGuard := context.WithTimeout(context.Background(), 5*time.Second) defer stopGuard() result := make(chan refreshRecommendScoresResult, 1) go func() { candidates, err := refreshRecommendScores( ctx, time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 1, 1, 1, store, ) result <- refreshRecommendScoresResult{candidates: candidates, err: err} }() select { case <-writeStarted: case <-guard.Done(): t.Fatalf("BulkWrite did not start: %v", guard.Err()) } cancel() var got refreshRecommendScoresResult select { case got = <-result: case <-guard.Done(): t.Fatalf("refreshRecommendScores() did not stop after cancellation: %v", guard.Err()) } if !errors.Is(got.err, context.Canceled) { t.Fatalf("refreshRecommendScores() error = %v, want %v", got.err, context.Canceled) } if got.candidates != nil { t.Fatalf("candidates = %v, want nil after cancellation", got.candidates) } select { case <-cancelObserved: default: t.Fatal("BulkWrite did not observe context cancellation") } } func TestRecommendScoreCursorQueryUsesBoundedIndexedScan(t *testing.T) { maxID := recommendBatchTestID(9) generatedAt := time.Date(2026, 7, 31, 8, 0, 0, 0, time.UTC) scope := newRecommendScoreScope(generatedAt, []string{" blocked ", "", "blocked", "other"}) filter, opts := recommendScoreCursorQuery(maxID, 4000, scope) if filter["status"] != CheckPass || filter["newsType"] != SHORT || filter["deleteAt"] != nil { t.Fatalf("filter = %#v, want eligible short-video filter", filter) } if !reflect.DeepEqual(filter["recoWeight"], bson.M{"$ne": -1}) { t.Fatalf("recoWeight filter = %#v, want explicit non-recommend exclusion", filter["recoWeight"]) } reviewAt, ok := filter["reviewAt"].(bson.M) if !ok || !reflect.DeepEqual(reviewAt, bson.M{"$lte": generatedAt}) { t.Fatalf("reviewAt filter = %#v, want <= %s", filter["reviewAt"], generatedAt) } moduleIDs, ok := filter["mId"].(bson.M) if !ok || !reflect.DeepEqual(moduleIDs["$nin"], []string{"blocked", "other"}) { t.Fatalf("mId filter = %#v, want normalized excluded modules", filter["mId"]) } idRange, ok := filter["_id"].(bson.M) if !ok || idRange["$lte"] != maxID || len(idRange) != 1 { t.Fatalf("_id filter = %#v, want $lte %s", filter["_id"], maxID.Hex()) } if opts.BatchSize == nil || *opts.BatchSize != 4000 { t.Fatalf("cursor batch size = %v, want 4000", opts.BatchSize) } if opts.Limit != nil { t.Fatalf("cursor limit = %v, want nil for one streaming cursor", *opts.Limit) } if opts.Hint != shortRecommendRefreshIndexName { t.Fatalf("cursor hint = %#v, want %s", opts.Hint, shortRecommendRefreshIndexName) } wantSort := bson.D{{Key: "_id", Value: 1}} if !reflect.DeepEqual(opts.Sort, wantSort) { t.Fatalf("cursor sort = %#v, want %#v", opts.Sort, wantSort) } projection, ok := opts.Projection.(bson.M) if !ok { t.Fatalf("cursor projection type = %T, want bson.M", opts.Projection) } for _, field := range []string{ "_id", "mId", "reviewAt", "likeCount", "collectCount", "commentCount", "shareCount", "recommendLikeCount", "recommendCollectCount", "recommendCommentCount", "recommendShareCount", "recommendScore", "recommendInitialized", } { if projection[field] != 1 { t.Fatalf("projection[%q] = %#v, want 1", field, projection[field]) } } } func TestRecommendScoreMaxEligibleIDQueryUsesSameEligibilityScope(t *testing.T) { generatedAt := time.Date(2026, 7, 31, 8, 0, 0, 0, time.UTC) scope := newRecommendScoreScope(generatedAt, []string{"blocked"}) filter, opts := recommendScoreMaxEligibleIDQuery(scope) if filter["status"] != CheckPass || filter["newsType"] != SHORT || filter["deleteAt"] != nil { t.Fatalf("filter = %#v, want eligible short-video filter", filter) } if !reflect.DeepEqual(filter["recoWeight"], bson.M{"$ne": -1}) { t.Fatalf("recoWeight filter = %#v, want explicit non-recommend exclusion", filter["recoWeight"]) } if !reflect.DeepEqual(filter["reviewAt"], bson.M{"$lte": generatedAt}) { t.Fatalf("reviewAt filter = %#v, want <= %s", filter["reviewAt"], generatedAt) } if !reflect.DeepEqual(filter["mId"], bson.M{"$nin": []string{"blocked"}}) { t.Fatalf("mId filter = %#v, want blocked module", filter["mId"]) } if opts.Hint != shortRecommendRefreshIndexName { t.Fatalf("hint = %#v, want %s", opts.Hint, shortRecommendRefreshIndexName) } if !reflect.DeepEqual(opts.Sort, bson.D{{Key: "_id", Value: -1}}) { t.Fatalf("sort = %#v, want descending _id", opts.Sort) } } func TestRecommendScoreMaxEligibleIDResultTreatsEmptyCollectionAsEmptyQueue(t *testing.T) { id, err := recommendScoreMaxEligibleIDResult( recommendBatchTestID(9), fmt.Errorf("wrapped: %w", mongo.ErrNoDocuments), ) if err != nil { t.Fatalf("recommendScoreMaxEligibleIDResult() error = %v", err) } if !id.IsZero() { t.Fatalf("id = %s, want zero ObjectID", id.Hex()) } wantErr := errors.New("find failed") id, err = recommendScoreMaxEligibleIDResult(recommendBatchTestID(9), wantErr) if !errors.Is(err, wantErr) || !id.IsZero() { t.Fatalf("id = %s, error = %v, want zero id and %v", id.Hex(), err, wantErr) } } func TestRecommendInitializationCountQueryUsesSameBoundedIndex(t *testing.T) { maxID := recommendBatchTestID(9) generatedAt := time.Date(2026, 7, 31, 8, 0, 0, 0, time.UTC) scope := newRecommendScoreScope(generatedAt, []string{"blocked"}) filter, opts := recommendInitializationCountQuery(maxID, 50_001, scope) if filter["status"] != CheckPass || filter["newsType"] != SHORT || filter["deleteAt"] != nil { t.Fatalf("filter = %#v, want eligible short-video filter", filter) } if !reflect.DeepEqual(filter["recoWeight"], bson.M{"$ne": -1}) { t.Fatalf("recoWeight filter = %#v, want explicit non-recommend exclusion", filter["recoWeight"]) } reviewAt, ok := filter["reviewAt"].(bson.M) if !ok || !reflect.DeepEqual(reviewAt, bson.M{"$lte": generatedAt}) { t.Fatalf("reviewAt filter = %#v, want <= %s", filter["reviewAt"], generatedAt) } moduleIDs, ok := filter["mId"].(bson.M) if !ok || !reflect.DeepEqual(moduleIDs["$nin"], []string{"blocked"}) { t.Fatalf("mId filter = %#v, want blocked module", filter["mId"]) } idRange, ok := filter["_id"].(bson.M) if !ok || idRange["$lte"] != maxID { t.Fatalf("_id filter = %#v", filter["_id"]) } initialized, ok := filter["recommendInitialized"].(bson.M) if !ok || initialized["$ne"] != true { t.Fatalf("recommendInitialized filter = %#v", filter["recommendInitialized"]) } if opts.Limit == nil || *opts.Limit != 50_001 { t.Fatalf("limit = %v, want 50001", opts.Limit) } if opts.Hint != shortRecommendRefreshIndexName { t.Fatalf("hint = %#v, want %s", opts.Hint, shortRecommendRefreshIndexName) } } func TestRecommendVideosByIDsFilterRevalidatesEligibility(t *testing.T) { ids := []ObjectID{recommendBatchTestID(1), recommendBatchTestID(2)} filter := recommendVideosByIDsFilter(ids) idMatch, ok := filter["_id"].(bson.M) if !ok || !reflect.DeepEqual(idMatch["$in"], ids) { t.Fatalf("_id filter = %#v", filter["_id"]) } if filter["status"] != CheckPass || filter["newsType"] != SHORT || filter["deleteAt"] != nil || !reflect.DeepEqual(filter["recoWeight"], bson.M{"$ne": -1}) { t.Fatalf("filter = %#v, want approved non-deleted short video", filter) } } func TestGetRecommendVideosByIDsContextEmptyInputDoesNotAccessMongo(t *testing.T) { originalDB := mdb mdb = nil t.Cleanup(func() { mdb = originalDB }) videos, err := GetRecommendVideosByIDsContext(context.Background(), nil) if err != nil { t.Fatalf("GetRecommendVideosByIDsContext() error = %v", err) } if len(videos) != 0 { t.Fatalf("videos = %#v, want empty", videos) } } func TestStreamRecommendScoreBatchesOwnsSlicesAndClosesCursor(t *testing.T) { documents := recommendBatchTestDocuments(5) maxID := documents[len(documents)-1].ID cursor := &fakeRecommendScoreCursor{documents: documents, current: -1} openCalls := 0 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return maxID, nil }, openCursor: func(_ context.Context, gotMaxID ObjectID, gotBatchSize int32) (recommendScoreCursor, error) { openCalls++ if gotMaxID != maxID { t.Fatalf("OpenCursor maxID = %s, want %s", gotMaxID.Hex(), maxID.Hex()) } if gotBatchSize != 4 { t.Fatalf("OpenCursor batch size = %d, want 4", gotBatchSize) } return cursor, nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return nil }, } var batches [][]recommendScoreDocument err := streamRecommendScoreBatches( context.Background(), maxID, 2, 4, store, func(batch []recommendScoreDocument) error { batches = append(batches, batch) return nil }, ) if err != nil { t.Fatalf("streamRecommendScoreBatches() error = %v", err) } if openCalls != 1 { t.Fatalf("OpenCursor calls = %d, want 1", openCalls) } if got, want := cursor.closeCalls.Load(), int32(1); got != want { t.Fatalf("Cursor close calls = %d, want %d", got, want) } if cursor.closeContextCanceled.Load() { t.Fatal("Cursor Close received an already-canceled cleanup context") } wantBatchSizes := []int{2, 2, 1} if len(batches) != len(wantBatchSizes) { t.Fatalf("batch count = %d, want %d", len(batches), len(wantBatchSizes)) } gotIDs := make([]ObjectID, 0, len(documents)) for i, batch := range batches { if len(batch) != wantBatchSizes[i] { t.Errorf("batch %d size = %d, want %d", i, len(batch), wantBatchSizes[i]) } for _, document := range batch { gotIDs = append(gotIDs, document.ID) } } for i := range documents { if gotIDs[i] != documents[i].ID { t.Fatalf("streamed ID %d = %s, want %s", i, gotIDs[i].Hex(), documents[i].ID.Hex()) } } } func TestStreamRecommendScoreBatchesErrorLifecycle(t *testing.T) { openErr := errors.New("open cursor failed") decodeErr := errors.New("decode failed") cursorErr := errors.New("cursor failed") closeErr := errors.New("close failed") tests := []struct { name string open func() (recommendScoreCursor, error) wantErr error wantClose int32 wantYieldCall int }{ { name: "open error", open: func() (recommendScoreCursor, error) { return nil, openErr }, wantErr: openErr, }, { name: "nil cursor", open: func() (recommendScoreCursor, error) { return nil, nil }, wantErr: errors.New("recommend score cursor is nil"), }, { name: "decode error wins over close error", open: func() (recommendScoreCursor, error) { return &fakeRecommendScoreCursor{ documents: recommendBatchTestDocuments(2), current: -1, decodeErrAt: 0, decodeErr: decodeErr, closeErr: closeErr, }, nil }, wantErr: decodeErr, wantClose: 1, }, { name: "cursor error prevents partial tail", open: func() (recommendScoreCursor, error) { return &fakeRecommendScoreCursor{ documents: recommendBatchTestDocuments(1), current: -1, cursorErr: cursorErr, }, nil }, wantErr: cursorErr, wantClose: 1, }, { name: "close error", open: func() (recommendScoreCursor, error) { return &fakeRecommendScoreCursor{ current: -1, closeErr: closeErr, }, nil }, wantErr: closeErr, wantClose: 1, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var opened recommendScoreCursor store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return recommendBatchTestID(2), nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { cursor, err := test.open() opened = cursor return cursor, err }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return nil }, } yieldCalls := 0 err := streamRecommendScoreBatches( context.Background(), recommendBatchTestID(2), 2, 4, store, func([]recommendScoreDocument) error { yieldCalls++ return nil }, ) if test.name == "nil cursor" { if err == nil || !strings.Contains(err.Error(), test.wantErr.Error()) { t.Fatalf("error = %v, want containing %q", err, test.wantErr) } } else if !errors.Is(err, test.wantErr) { t.Fatalf("error = %v, want %v", err, test.wantErr) } if yieldCalls != test.wantYieldCall { t.Fatalf("yield calls = %d, want %d", yieldCalls, test.wantYieldCall) } if cursor, ok := opened.(*fakeRecommendScoreCursor); ok { if got := cursor.closeCalls.Load(); got != test.wantClose { t.Fatalf("Cursor close calls = %d, want %d", got, test.wantClose) } } }) } } func TestStreamRecommendScoreBatchesCancellationUsesCleanupContext(t *testing.T) { documents := recommendBatchTestDocuments(2) cursor := &fakeRecommendScoreCursor{documents: documents, current: -1} store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return documents[len(documents)-1].ID, nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { return cursor, nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return nil }, } ctx, cancel := context.WithCancel(context.Background()) err := streamRecommendScoreBatches( ctx, documents[len(documents)-1].ID, 1, 2, store, func([]recommendScoreDocument) error { cancel() return context.Canceled }, ) if !errors.Is(err, context.Canceled) { t.Fatalf("error = %v, want %v", err, context.Canceled) } if got := cursor.closeCalls.Load(); got != 1 { t.Fatalf("Cursor close calls = %d, want 1", got) } if cursor.closeContextCanceled.Load() { t.Fatal("Cursor Close received the canceled work context") } } func TestStreamRecommendScoreBatchesRejectsInvalidIDSequence(t *testing.T) { id1 := recommendBatchTestID(1) id2 := recommendBatchTestID(2) id3 := recommendBatchTestID(3) tests := []struct { name string documents []recommendScoreDocument maxID ObjectID }{ { name: "zero ID", documents: []recommendScoreDocument{{}}, maxID: id2, }, { name: "duplicate across batch boundary", documents: []recommendScoreDocument{{ID: id1}, {ID: id1}}, maxID: id2, }, { name: "descending across batch boundary", documents: []recommendScoreDocument{{ID: id2}, {ID: id1}}, maxID: id2, }, { name: "above fixed high water", documents: []recommendScoreDocument{{ID: id1}, {ID: id3}}, maxID: id2, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { cursor := &fakeRecommendScoreCursor{documents: test.documents, current: -1} store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return test.maxID, nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { return cursor, nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return nil }, } err := streamRecommendScoreBatches( context.Background(), test.maxID, 1, 2, store, func([]recommendScoreDocument) error { return nil }, ) if err == nil || !strings.Contains(err.Error(), "invalid _id sequence") { t.Fatalf("error = %v, want invalid _id sequence", err) } if got := cursor.closeCalls.Load(); got != 1 { t.Fatalf("Cursor close calls = %d, want 1", got) } }) } } func TestRefreshRecommendScoresNonZeroHighWaterEmptyCursor(t *testing.T) { highWater := recommendBatchTestID(1) cursor := &fakeRecommendScoreCursor{current: -1} var bulkCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { return highWater, nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { return cursor, nil }, bulkWrite: func(context.Context, []mongo.WriteModel) error { bulkCalls.Add(1) return nil }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 4, 2, store, ) if err != nil { t.Fatalf("refreshRecommendScores() error = %v", err) } if candidates == nil || len(candidates) != 0 { t.Fatalf("candidates = %#v, want non-nil empty slice", candidates) } if got := cursor.closeCalls.Load(); got != 1 { t.Fatalf("Cursor close calls = %d, want 1", got) } if got := bulkCalls.Load(); got != 0 { t.Fatalf("BulkWrite calls = %d, want 0", got) } } func TestRefreshRecommendScoresRejectsInvalidLimitsBeforeDatabaseAccess(t *testing.T) { tests := []struct { name string batchSize int cursorBatchSize int32 workerCount int wantMessage string }{ { name: "process batch", batchSize: 0, cursorBatchSize: 4, workerCount: 2, wantMessage: "batch size must be positive", }, { name: "cursor batch", batchSize: 2, cursorBatchSize: 0, workerCount: 2, wantMessage: "cursor batch size must be positive", }, { name: "worker count", batchSize: 2, cursorBatchSize: 4, workerCount: 0, wantMessage: "worker count must be positive", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var databaseCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: func(context.Context) (ObjectID, error) { databaseCalls.Add(1) return recommendBatchTestID(1), nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { databaseCalls.Add(1) return nil, errors.New("must not open cursor") }, bulkWrite: func(context.Context, []mongo.WriteModel) error { databaseCalls.Add(1) return nil }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), test.batchSize, test.cursorBatchSize, test.workerCount, store, ) if err == nil || !strings.Contains(err.Error(), test.wantMessage) { t.Fatalf("error = %v, want containing %q", err, test.wantMessage) } if candidates != nil { t.Fatalf("candidates = %#v, want nil", candidates) } if got := databaseCalls.Load(); got != 0 { t.Fatalf("database calls = %d, want 0", got) } }) } } func TestRefreshRecommendScoresPropagatesSourceErrors(t *testing.T) { maxErr := errors.New("max ID failed") openErr := errors.New("open cursor failed") tests := []struct { name string maxID func(context.Context) (ObjectID, error) openCursor func(context.Context, ObjectID, int32) (recommendScoreCursor, error) wantErr error wantOpens int32 }{ { name: "max ID", maxID: func(context.Context) (ObjectID, error) { return ObjectID{}, maxErr }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { return nil, errors.New("must not open cursor") }, wantErr: maxErr, }, { name: "open cursor", maxID: func(context.Context) (ObjectID, error) { return recommendBatchTestID(1), nil }, openCursor: func(context.Context, ObjectID, int32) (recommendScoreCursor, error) { return nil, openErr }, wantErr: openErr, wantOpens: 1, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var openCalls atomic.Int32 store := fakeRecommendScoreStore{ maxEligibleID: test.maxID, openCursor: func(ctx context.Context, maxID ObjectID, batchSize int32) (recommendScoreCursor, error) { openCalls.Add(1) return test.openCursor(ctx, maxID, batchSize) }, bulkWrite: func(context.Context, []mongo.WriteModel) error { return errors.New("must not write") }, } candidates, err := refreshRecommendScores( context.Background(), time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC), 2, 4, 2, store, ) if !errors.Is(err, test.wantErr) { t.Fatalf("error = %v, want %v", err, test.wantErr) } if candidates != nil { t.Fatalf("candidates = %#v, want nil", candidates) } if got := openCalls.Load(); got != test.wantOpens { t.Fatalf("OpenCursor calls = %d, want %d", got, test.wantOpens) } }) } }