1492 lines
42 KiB
Go
1492 lines
42 KiB
Go
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)
|
|
}
|
|
})
|
|
}
|
|
}
|