iotlcm/admin-svr/server/internal/logic/common/batch_main_requeue_test.go

package common

import (
	"context"
	"errors"
	"testing"
	"time"

	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/consts"
	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/model"

	_ "github.com/gogf/gf/contrib/drivers/sqlite/v2"
	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gtime"
)

func TestBatchMainRequeuePublishesPendingBatch(t *testing.T) {
	ctx := setupBatchMainRequeueDB(t)
	insertRequeueBatch(t, ctx, "PENDING-1", consts.BatchTypeInbound, consts.BatchAsyncStatusPending, consts.BatchRunResultPending, `{}`)

	var published []batchExecuteMessage
	restore := replaceBatchRequeuePublisher(func(_ context.Context, queue string, data any) error {
		if queue != consts.QueueBatchExecute {
			t.Fatalf("queue = %s, want %s", queue, consts.QueueBatchExecute)
		}
		message, ok := data.(batchExecuteMessage)
		if !ok {
			t.Fatalf("message type = %T, want batchExecuteMessage", data)
		}
		published = append(published, message)
		return nil
	})
	defer restore()

	res, err := NewCommonBatchMain().Requeue(ctx, &model.BatchMainRequeueParam{BatchNo: "PENDING-1", Reason: "message lost"})
	if err != nil {
		t.Fatalf("Requeue() error = %v", err)
	}
	if res == nil || !res.Requeued || res.BatchNo != "PENDING-1" || res.BatchType != consts.BatchTypeInbound {
		t.Fatalf("unexpected response: %+v", res)
	}
	if len(published) != 1 || published[0].BatchNo != "PENDING-1" || published[0].BatchType != consts.BatchTypeInbound {
		t.Fatalf("unexpected published messages: %+v", published)
	}
}

func TestBatchMainRequeuePublishesDueRestartWaitingResume(t *testing.T) {
	ctx := setupBatchMainRequeueDB(t)
	resumeAt := gtime.Now().Add(-time.Minute).String()
	insertRequeueBatch(t, ctx, "RST-DUE", consts.BatchTypeRestart, consts.BatchAsyncStatusProcessing, consts.BatchRunResultRunning, `{"restart":{"phase":"waiting_resume","resumeAt":"`+resumeAt+`"}}`)

	var published []batchExecuteMessage
	restore := replaceBatchRequeuePublisher(func(_ context.Context, _ string, data any) error {
		published = append(published, data.(batchExecuteMessage))
		return nil
	})
	defer restore()

	res, err := NewCommonBatchMain().Requeue(ctx, &model.BatchMainRequeueParam{BatchNo: "RST-DUE", Reason: "delayed message lost"})
	if err != nil {
		t.Fatalf("Requeue() error = %v", err)
	}
	if res == nil || res.BatchNo != "RST-DUE" || res.BatchType != consts.BatchTypeRestart {
		t.Fatalf("unexpected response: %+v", res)
	}
	if len(published) != 1 || published[0].BatchNo != "RST-DUE" || published[0].BatchType != consts.BatchTypeRestart {
		t.Fatalf("unexpected published messages: %+v", published)
	}
}

func TestBatchMainRequeueRejectsUnsafeStatuses(t *testing.T) {
	ctx := setupBatchMainRequeueDB(t)
	future := gtime.Now().Add(time.Hour).String()
	cases := []struct {
		name      string
		batchNo   string
		batchType int
		status    int
		runResult int
		info      string
	}{
		{name: "failed", batchNo: "FAILED-1", batchType: consts.BatchTypeInbound, status: consts.BatchAsyncStatusFailed, runResult: consts.BatchRunResultFailed, info: `{}`},
		{name: "success", batchNo: "SUCCESS-1", batchType: consts.BatchTypeInbound, status: consts.BatchAsyncStatusSuccess, runResult: consts.BatchRunResultSuccess, info: `{}`},
		{name: "partial", batchNo: "PARTIAL-1", batchType: consts.BatchTypeInbound, status: consts.BatchAsyncStatusPartialSuccess, runResult: consts.BatchRunResultPartialSuccess, info: `{}`},
		{name: "review", batchNo: "REVIEW-1", batchType: consts.BatchTypeDataStart, status: consts.BatchAsyncStatusReview, runResult: consts.BatchRunResultReview, info: `{}`},
		{name: "normal processing", batchNo: "PROCESSING-1", batchType: consts.BatchTypeInbound, status: consts.BatchAsyncStatusProcessing, runResult: consts.BatchRunResultRunning, info: `{}`},
		{name: "restart future resume", batchNo: "RST-FUTURE", batchType: consts.BatchTypeRestart, status: consts.BatchAsyncStatusProcessing, runResult: consts.BatchRunResultRunning, info: `{"restart":{"phase":"waiting_resume","resumeAt":"` + future + `"}}`},
	}
	for _, tc := range cases {
		insertRequeueBatch(t, ctx, tc.batchNo, tc.batchType, tc.status, tc.runResult, tc.info)
	}

	var published int
	restore := replaceBatchRequeuePublisher(func(context.Context, string, any) error {
		published++
		return nil
	})
	defer restore()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := NewCommonBatchMain().Requeue(ctx, &model.BatchMainRequeueParam{BatchNo: tc.batchNo, Reason: "operator retry"})
			if err == nil {
				t.Fatal("Requeue() error = nil, want error")
			}
		})
	}
	if published != 0 {
		t.Fatalf("published messages = %d, want 0", published)
	}
}

func TestBatchMainRequeueHandlesMissingBatchAndPublishFailure(t *testing.T) {
	ctx := setupBatchMainRequeueDB(t)
	insertRequeueBatch(t, ctx, "PENDING-FAIL-PUBLISH", consts.BatchTypeInbound, consts.BatchAsyncStatusPending, consts.BatchRunResultPending, `{}`)

	restore := replaceBatchRequeuePublisher(func(context.Context, string, any) error {
		return errors.New("rabbitmq unavailable")
	})
	defer restore()

	if _, err := NewCommonBatchMain().Requeue(ctx, &model.BatchMainRequeueParam{BatchNo: "PENDING-FAIL-PUBLISH", Reason: "message lost"}); err == nil {
		t.Fatal("Requeue() publish failure error = nil, want error")
	}
	if _, err := NewCommonBatchMain().Requeue(ctx, &model.BatchMainRequeueParam{BatchNo: "MISSING", Reason: "message lost"}); err == nil {
		t.Fatal("Requeue() missing batch error = nil, want error")
	}
}

func setupBatchMainRequeueDB(t *testing.T) context.Context {
	t.Helper()
	ctx := context.Background()
	if err := gdb.SetConfig(gdb.Config{
		"default": gdb.ConfigGroup{{
			Type:             "sqlite",
			Link:             "sqlite::@file(:memory:)",
			MaxOpenConnCount: 1,
		}},
	}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		_ = g.DB().Close(ctx)
	})
	// GoFrame caches table fields by group and table, across SQLite instances.
	if err := g.DB().GetCore().ClearTableFieldsAll(ctx); err != nil {
		t.Fatal(err)
	}
	if _, err := g.DB().Exec(ctx, `CREATE TABLE batch_main (
            source_channel TEXT NOT NULL DEFAULT 'legacy',
            execution_mode INTEGER NOT NULL DEFAULT 0,
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		batch_no TEXT NOT NULL,
		batch_type INTEGER NOT NULL,
		batch_async_status INTEGER NOT NULL,
		run_result INTEGER NOT NULL,
		batch_info TEXT,
		tenant_id TEXT NOT NULL DEFAULT '',
		dept_id INTEGER NOT NULL DEFAULT 0,
		created_dept INTEGER NOT NULL DEFAULT 0,
		created_by INTEGER NOT NULL DEFAULT 0
	)`); err != nil {
		t.Fatal(err)
	}
	return ctx
}

func insertRequeueBatch(t *testing.T, ctx context.Context, batchNo string, batchType, status, runResult int, batchInfo string) {
	t.Helper()
	if _, err := g.DB().Exec(ctx,
		`INSERT INTO batch_main(batch_no, batch_type, batch_async_status, run_result, batch_info, tenant_id) VALUES(?, ?, ?, ?, ?, '')`,
		batchNo, batchType, status, runResult, batchInfo,
	); err != nil {
		t.Fatal(err)
	}
}

func replaceBatchRequeuePublisher(fn func(context.Context, string, any) error) func() {
	original := publishBatchExecuteMessage
	publishBatchExecuteMessage = fn
	return func() {
		publishBatchExecuteMessage = original
	}
}