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

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)
}