chore: 并入 iotlcm-worker-svr 历史到 worker-svr/（原提交 96f0f4e）

worker-svr/internal/producer/monthlyusage/producer_test.go

@@ -0,0 +1,168 @@
+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)
+}