chore: 并入 iotlcm-worker-svr 历史到 worker-svr/（原提交 96f0f4e）
This commit is contained in:
MiFaZhan
@@ -0,0 +1,274 @@
|
||||
package zset
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gogf/gf/v2/container/gvar"
|
||||
)
|
||||
|
||||
type redisCall struct {
|
||||
command string
|
||||
args []any
|
||||
}
|
||||
|
||||
type fakeRedis struct {
|
||||
calls []redisCall
|
||||
responses []*gvar.Var
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeRedis) Do(_ context.Context, command string, args ...any) (*gvar.Var, error) {
|
||||
f.calls = append(f.calls, redisCall{command: command, args: args})
|
||||
if f.err != nil {
|
||||
return nil, f.err
|
||||
}
|
||||
if len(f.responses) == 0 {
|
||||
return gvar.New(int64(1)), nil
|
||||
}
|
||||
response := f.responses[0]
|
||||
f.responses = f.responses[1:]
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (f *fakeRedis) Eval(_ context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
||||
allArgs := []any{script, numKeys}
|
||||
for _, k := range keys {
|
||||
allArgs = append(allArgs, k)
|
||||
}
|
||||
allArgs = append(allArgs, args...)
|
||||
f.calls = append(f.calls, redisCall{command: "EVAL", args: allArgs})
|
||||
if f.err != nil {
|
||||
return nil, f.err
|
||||
}
|
||||
if len(f.responses) == 0 {
|
||||
return gvar.New(int64(1)), nil
|
||||
}
|
||||
response := f.responses[0]
|
||||
f.responses = f.responses[1:]
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func TestStoreUsesContractKeys(t *testing.T) {
|
||||
store, err := New(&fakeRedis{}, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key, err := store.Key("mobile-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if key != "flow:mobile-a" {
|
||||
t.Fatalf("task key = %q, want flow:mobile-a", key)
|
||||
}
|
||||
lease, err := store.leaseKey("mobile-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if lease != "flow:lease:mobile-a" {
|
||||
t.Fatalf("lease key = %q, want flow:lease:mobile-a", lease)
|
||||
}
|
||||
if got := store.ChannelsKey(); got != "flow:channels" {
|
||||
t.Fatalf("channels key = %q, want flow:channels", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStableTaskID(t *testing.T) {
|
||||
if got := StableTaskID("run-1", "current_month_usage", 42); got != "run-1:current_month_usage:42" {
|
||||
t.Fatalf("stable task id = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunKeysShareRedisHashTag(t *testing.T) {
|
||||
store, err := New(&fakeRedis{}, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys := []func() (string, error){
|
||||
func() (string, error) { return store.RunMetaKey("run-1") },
|
||||
func() (string, error) { return store.RoundActiveKey("run-1", 1, "mobile-a") },
|
||||
func() (string, error) { return store.RoundWaitingKey("run-1", 2, "mobile-a") },
|
||||
func() (string, error) { return store.RoundLeaseKey("run-1", 1, "mobile-a") },
|
||||
func() (string, error) { return store.TaskStateKey("run-1") },
|
||||
func() (string, error) { return store.FailedKey("run-1") },
|
||||
}
|
||||
for _, build := range keys {
|
||||
key, err := build()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(key, "flow:{run-1}") {
|
||||
t.Fatalf("run key %q does not contain hash tag", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnqueueMemberExcludesSchedulingScore(t *testing.T) {
|
||||
fake := &fakeRedis{responses: []*gvar.Var{gvar.New(int64(1))}}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task := Task{TaskID: "run:current_month_usage:1", ChanKey: "mobile-a", CardID: 1}
|
||||
if err = store.Enqueue(context.Background(), task, time.UnixMilli(1234)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(fake.calls) != 1 || fake.calls[0].command != "EVAL" {
|
||||
t.Fatalf("Redis calls = %+v", fake.calls)
|
||||
}
|
||||
member, ok := fake.calls[0].args[len(fake.calls[0].args)-1].(string)
|
||||
if !ok {
|
||||
t.Fatalf("member type = %T", fake.calls[0].args[len(fake.calls[0].args)-1])
|
||||
}
|
||||
if strings.Contains(member, "nextExecuteAt") {
|
||||
t.Fatalf("member contains volatile score: %s", member)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimDecodesMemberAndUsesLeaseKey(t *testing.T) {
|
||||
task := Task{
|
||||
TaskID: "run-1:current_month_usage:42",
|
||||
RunID: "run-1",
|
||||
Operation: "current_month_usage",
|
||||
QueryMonth: "202608",
|
||||
ChanKey: "mobile-a",
|
||||
CardID: 42,
|
||||
ICCID: "8986000000000042",
|
||||
}
|
||||
member, err := json.Marshal(task)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fake := &fakeRedis{responses: []*gvar.Var{gvar.New([]any{string(member), "worker:token"})}}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.UnixMilli(1000)
|
||||
leaseUntil := now.Add(time.Minute)
|
||||
claimed, err := store.Claim(context.Background(), "mobile-a", now, leaseUntil, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(claimed) != 1 || claimed[0].Task.TaskID != task.TaskID {
|
||||
t.Fatalf("claimed tasks = %+v", claimed)
|
||||
}
|
||||
if claimed[0].LeaseToken != "worker:token" {
|
||||
t.Fatalf("lease token = %q", claimed[0].LeaseToken)
|
||||
}
|
||||
if len(fake.calls) != 1 || fake.calls[0].command != "EVAL" {
|
||||
t.Fatalf("Redis calls = %+v", fake.calls)
|
||||
}
|
||||
args := fake.calls[0].args
|
||||
if len(args) < 4 || args[3] != "flow:lease:mobile-a" {
|
||||
t.Fatalf("claim args do not use lease key: %+v", args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAckBatchUsesLuaOwnershipProtocol(t *testing.T) {
|
||||
fake := &fakeRedis{}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
claimed := ClaimedTask{
|
||||
Task: Task{TaskID: "run:current_month_usage:1", ChanKey: "mobile-a", CardID: 1},
|
||||
Member: `{"taskId":"run:current_month_usage:1","chanKey":"mobile-a","cardId":1}`,
|
||||
LeaseToken: "token-1",
|
||||
}
|
||||
count, err := store.AckBatch(context.Background(), []ClaimedTask{claimed})
|
||||
if err != nil || count != 1 {
|
||||
t.Fatalf("ack count = %d, err = %v", count, err)
|
||||
}
|
||||
renewed, err := store.Renew(context.Background(), claimed, time.Now().Add(time.Minute))
|
||||
if err != nil || !renewed {
|
||||
t.Fatalf("renew = %t, err = %v", renewed, err)
|
||||
}
|
||||
if len(fake.calls) != 2 {
|
||||
t.Fatalf("Redis calls = %d, want 2", len(fake.calls))
|
||||
}
|
||||
for _, call := range fake.calls {
|
||||
if call.command != "EVAL" {
|
||||
t.Fatalf("command = %q, want EVAL", call.command)
|
||||
}
|
||||
script, ok := call.args[0].(string)
|
||||
if !ok || !strings.Contains(script, "HGET") {
|
||||
t.Fatalf("Lua script does not validate ownership: %v", call.args[0])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenewBatchUsesOneAtomicLuaCall(t *testing.T) {
|
||||
fake := &fakeRedis{responses: []*gvar.Var{gvar.New(int64(2))}}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tasks := []ClaimedTask{
|
||||
{Task: Task{ChanKey: "mobile-a"}, Member: "member-1", LeaseToken: "token-1"},
|
||||
{Task: Task{ChanKey: "mobile-a"}, Member: "member-2", LeaseToken: "token-2"},
|
||||
}
|
||||
count, err := store.RenewBatch(context.Background(), tasks, time.UnixMilli(5000))
|
||||
if err != nil || count != 2 {
|
||||
t.Fatalf("RenewBatch() = %d, err = %v", count, err)
|
||||
}
|
||||
if len(fake.calls) != 1 || fake.calls[0].command != "EVAL" {
|
||||
t.Fatalf("Redis calls = %+v, want one EVAL", fake.calls)
|
||||
}
|
||||
script := fake.calls[0].args[0].(string)
|
||||
if !strings.Contains(script, "HGET") || !strings.Contains(script, "ZSCORE") {
|
||||
t.Fatalf("renew script does not validate all leases: %s", script)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnqueueBatchRegistersChannelAndWritesTasksOnce(t *testing.T) {
|
||||
fake := &fakeRedis{responses: []*gvar.Var{gvar.New(int64(2))}}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tasks := []Task{
|
||||
{TaskID: "run:usage:1", RunID: "run", ChanKey: "mobile-a", CardID: 1},
|
||||
{TaskID: "run:usage:2", RunID: "run", ChanKey: "mobile-a", CardID: 2},
|
||||
}
|
||||
added, err := store.EnqueueBatch(context.Background(), "mobile-a", tasks, time.UnixMilli(1234))
|
||||
if err != nil || added != 2 {
|
||||
t.Fatalf("EnqueueBatch() = %d, err = %v", added, err)
|
||||
}
|
||||
if len(fake.calls) != 1 || fake.calls[0].command != "EVAL" {
|
||||
t.Fatalf("Redis calls = %+v, want one EVAL", fake.calls)
|
||||
}
|
||||
if fake.calls[0].args[2] != "flow:mobile-a" || fake.calls[0].args[3] != "flow:channels" || fake.calls[0].args[4] != "flow:blocked:mobile-a" {
|
||||
t.Fatalf("batch keys = %+v", fake.calls[0].args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnqueueBatchRejectsMixedChannels(t *testing.T) {
|
||||
store, err := New(&fakeRedis{}, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = store.EnqueueBatch(context.Background(), "mobile-a", []Task{{TaskID: "1", RunID: "run", ChanKey: "mobile-b"}}, time.Now())
|
||||
if err == nil {
|
||||
t.Fatal("expected mixed channel batch error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlockRunUsesTaskLeaseAndBlockedKeys(t *testing.T) {
|
||||
fake := &fakeRedis{responses: []*gvar.Var{gvar.New(int64(7))}}
|
||||
store, err := New(fake, "flow:")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
removed, err := store.BlockRun(context.Background(), "telecom-a", "run-1")
|
||||
if err != nil || removed != 7 {
|
||||
t.Fatalf("BlockRun() = %d, err = %v", removed, err)
|
||||
}
|
||||
args := fake.calls[0].args
|
||||
if args[2] != "flow:telecom-a" || args[3] != "flow:lease:telecom-a" || args[4] != "flow:blocked:telecom-a" {
|
||||
t.Fatalf("block keys = %+v", args)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user