107 lines
3.4 KiB
Go
107 lines
3.4 KiB
Go
package session
|
||
|
||
import (
|
||
"context"
|
||
"os"
|
||
"testing"
|
||
"time"
|
||
|
||
"ai-agent/workflow/consts/flow"
|
||
sessionDto "ai-agent/workflow/model/dto/session"
|
||
"ai-agent/workflow/model/entity"
|
||
|
||
"gitea.redpowerfuture.com/red-future/common/beans"
|
||
_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
|
||
)
|
||
|
||
// TestExecWorkflowRetryColsRoundTrip 验证三新列读写 + ListRecoverable 三类判定（需 AI_AGENT_TEST_DB=1 且本地 PG 可用）
|
||
func TestExecWorkflowRetryColsRoundTrip(t *testing.T) {
|
||
if !isTestDBEnabled(t) {
|
||
t.Skip("skip: AI_AGENT_TEST_DB not set")
|
||
}
|
||
// gfdb 的 Hook 依赖 ctx 中注入用户信息（租户/创建人），否则 Insert 报 "token 数据为空"
|
||
ctx := context.WithValue(context.Background(), "user", &beans.User{
|
||
UserName: "test-retry", TenantId: 1,
|
||
})
|
||
id, err := ExecWorkflowDao.Insert(ctx, &sessionDto.CreateWorkflowReq{
|
||
SessionId: "test-retry", FlowId: 1, NodeGroupId: "ng",
|
||
Status: flow.FlowExecutionStatusRunning.Code(),
|
||
LastHeartbeat: time.Now().UnixMilli(),
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("insert: %v", err)
|
||
}
|
||
defer ExecWorkflowDao.Delete(ctx, &sessionDto.DeleteExecWorkflowReq{Id: []int64{id}})
|
||
|
||
if err := ExecWorkflowDao.TouchHeartbeat(ctx, id); err != nil {
|
||
t.Fatalf("touch heartbeat: %v", err)
|
||
}
|
||
if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 1); err != nil {
|
||
t.Fatalf("update retry: %v", err)
|
||
}
|
||
got, err := ExecWorkflowDao.GetById(ctx, id)
|
||
if err != nil || got == nil {
|
||
t.Fatalf("get: %v", err)
|
||
}
|
||
if got.Retryable != 1 || got.RetryCount != 1 || got.LastHeartbeat == 0 {
|
||
t.Fatalf("cols not persisted: retryable=%d retryCount=%d heartbeat=%d", got.Retryable, got.RetryCount, got.LastHeartbeat)
|
||
}
|
||
|
||
now := time.Now().UnixMilli()
|
||
stale := now - int64(120*time.Second/time.Millisecond)
|
||
if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
|
||
t.Fatalf("update retry reset: %v", err)
|
||
}
|
||
// 情况1：status=3 retryable=1 retry_count<N → 可恢复
|
||
_, err = ExecWorkflowDao.Update(ctx, &sessionDto.UpdateWorkflowReq{Id: id, Status: flow.FlowExecutionStatusFailed.Code()})
|
||
if err != nil {
|
||
t.Fatalf("update failed: %v", err)
|
||
}
|
||
if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 0); err != nil {
|
||
t.Fatalf("update retry: %v", err)
|
||
}
|
||
rows, err := ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
|
||
if err != nil {
|
||
t.Fatalf("list recoverable: %v", err)
|
||
}
|
||
if !containsID(rows, id) {
|
||
t.Fatalf("status=3 retryable=1 应可恢复，未命中 id=%d", id)
|
||
}
|
||
// 情况2：retry_count 耗尽 → 不可恢复
|
||
if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 2); err != nil {
|
||
t.Fatalf("update retry exhaust: %v", err)
|
||
}
|
||
rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
|
||
if err != nil {
|
||
t.Fatalf("list recoverable: %v", err)
|
||
}
|
||
if containsID(rows, id) {
|
||
t.Fatalf("retry 耗尽不应可恢复")
|
||
}
|
||
// 情况3：status=3 retryable=0（用户取消）→ 不可恢复
|
||
if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
|
||
t.Fatalf("update retry: %v", err)
|
||
}
|
||
rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
|
||
if err != nil {
|
||
t.Fatalf("list recoverable: %v", err)
|
||
}
|
||
if containsID(rows, id) {
|
||
t.Fatalf("用户取消不可恢复")
|
||
}
|
||
}
|
||
|
||
func containsID(rows []*entity.ExecWorkflow, id int64) bool {
|
||
for _, r := range rows {
|
||
if r.Id == id {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// isTestDBEnabled 测试辅助：AI_AGENT_TEST_DB=1 时连接真实 PG
|
||
func isTestDBEnabled(t *testing.T) bool {
|
||
return os.Getenv("AI_AGENT_TEST_DB") == "1"
|
||
}
|