package monthlyusage import ( "context" "fmt" "strings" "sync" "testing" "github.com/gogf/gf/v2/container/gvar" "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/queue/zset" ) type pendingRedis struct { channels []string sizes map[string]int64 failKey string } func (r *pendingRedis) Do(_ context.Context, command string, args ...any) (*gvar.Var, error) { switch strings.ToUpper(command) { case "SMEMBERS": return gvar.New(r.channels), nil case "ZCARD": key := fmt.Sprint(args[0]) if key == r.failKey { return nil, fmt.Errorf("redis unavailable") } return gvar.New(r.sizes[key]), nil default: return nil, fmt.Errorf("unexpected command %s", command) } } func (r *pendingRedis) Eval(_ context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) { return nil, fmt.Errorf("unexpected Eval call") } func TestHasPendingTasks(t *testing.T) { tests := []struct { name string redis *pendingRedis want bool wantErr bool }{ { name: "all registered channels are empty", redis: &pendingRedis{ channels: []string{"mobile-a", "telecom-a"}, sizes: map[string]int64{ "flow:mobile-a": 0, "flow:telecom-a": 0, }, }, }, { name: "previous run still has tasks", redis: &pendingRedis{ channels: []string{"mobile-a", "telecom-a"}, sizes: map[string]int64{ "flow:mobile-a": 0, "flow:telecom-a": 3, }, }, want: true, }, { name: "redis error is not treated as empty", redis: &pendingRedis{ channels: []string{"mobile-a"}, sizes: map[string]int64{}, failKey: "flow:mobile-a", }, wantErr: true, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { store, err := zset.New(test.redis, "flow:") if err != nil { t.Fatal(err) } producer, err := New(&fakeScanner{}, store, Config{}) if err != nil { t.Fatal(err) } got, err := func() (bool, error) { producer.config.RoundRetryEnabled = false return producer.hasPendingTasks(context.Background()) }() if (err != nil) != test.wantErr { t.Fatalf("hasPendingTasks() error = %v, wantErr %v", err, test.wantErr) } if got != test.want { t.Fatalf("hasPendingTasks() = %v, want %v", got, test.want) } }) } } type fakeScanner struct{} func (*fakeScanner) SnapshotMaxID(context.Context) (int64, error) { return 0, nil } func (*fakeScanner) Scan(context.Context, int64, int64, int) ([]Card, error) { return nil, nil } func (*fakeScanner) ListChannels(context.Context) ([]channelMetadata, error) { return nil, nil } func (*fakeScanner) ScanChannel(context.Context, string, int64, int) ([]Card, error) { return nil, nil } type pageScanner struct { cards []Card } func (s *pageScanner) SnapshotMaxID(context.Context) (int64, error) { if len(s.cards) == 0 { return 0, nil } return s.cards[len(s.cards)-1].ID, nil } func (s *pageScanner) Scan(_ context.Context, lastID, _ int64, _ int) ([]Card, error) { if lastID != 0 { return nil, nil } return s.cards, nil } type producerRedis struct { mu sync.Mutex enqueueBatches map[string]int } func (r *producerRedis) Do(_ context.Context, command string, _ ...any) (*gvar.Var, error) { switch strings.ToUpper(command) { case "SET": return gvar.New("OK"), nil case "SMEMBERS": return gvar.New([]string{}), nil case "HGETALL": // No stored run metadata: the producer treats this as a first delivery. return gvar.New(map[string]string{}), nil default: return nil, fmt.Errorf("unexpected command %s", command) } } func (r *producerRedis) Eval(_ context.Context, _ string, numKeys int64, keys []string, args []any) (*gvar.Var, error) { if numKeys == 1 { return gvar.New(int64(1)), nil } if numKeys != 3 { return nil, fmt.Errorf("unexpected key count %d", numKeys) } r.mu.Lock() r.enqueueBatches[keys[0]] = len(args) - 3 r.mu.Unlock() return gvar.New(len(args) - 3), nil } func (s *pageScanner) ListChannels(ctx context.Context) ([]channelMetadata, error) { return scannerFor(s.cards, 0).ListChannels(ctx) } func (s *pageScanner) ScanChannel(ctx context.Context, key string, last int64, size int) ([]Card, error) { return scannerFor(s.cards, 0).ScanChannel(ctx, key, last, size) }