61 lines
1.8 KiB
Go
61 lines
1.8 KiB
Go
package flow
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"ai-agent/workflow/consts/flow"
|
|
"ai-agent/workflow/model/entity"
|
|
)
|
|
|
|
func TestIsRecoverable(t *testing.T) {
|
|
now := time.Now().UnixMilli()
|
|
staleBefore := now - int64(heartbeatStaleAfter/time.Millisecond) - 1
|
|
fresh := now // 心跳新鲜
|
|
|
|
running := func(hb int64) *entity.ExecWorkflow {
|
|
return &entity.ExecWorkflow{Status: flow.FlowExecutionStatusRunning.Code(), LastHeartbeat: hb}
|
|
}
|
|
failed := func(retryable, retryCount int) *entity.ExecWorkflow {
|
|
return &entity.ExecWorkflow{Status: flow.FlowExecutionStatusFailed.Code(), Retryable: retryable, RetryCount: retryCount}
|
|
}
|
|
|
|
cases := []struct {
|
|
name string
|
|
exec *entity.ExecWorkflow
|
|
want bool
|
|
}{
|
|
{"nil", nil, false},
|
|
{"status=1 心跳陈旧 → 恢复", running(staleBefore), true},
|
|
{"status=1 心跳新鲜 → 不恢复", running(fresh), false},
|
|
{"status=1 心跳为0（老数据）→ 恢复", running(0), true},
|
|
{"status=3 retryable=1 未耗尽 → 恢复", failed(1, 0), true},
|
|
{"status=3 retryable=1 已耗尽 → 不恢复", failed(1, execMaxRetryCount), false},
|
|
{"status=3 retryable=0（用户取消）→ 不恢复", failed(0, 0), false},
|
|
{"status=2 成功 → 不恢复", &entity.ExecWorkflow{Status: flow.FlowExecutionStatusSuccess.Code()}, false},
|
|
}
|
|
for _, c := range cases {
|
|
if got := isRecoverable(c.exec, now); got != c.want {
|
|
t.Fatalf("%s: got %v want %v", c.name, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestShouldRetry(t *testing.T) {
|
|
canceled := context.Canceled
|
|
if shouldRetry(canceled) {
|
|
t.Fatalf("用户取消不应重试")
|
|
}
|
|
if shouldRetry(nil) {
|
|
t.Fatalf("nil 不应重试")
|
|
}
|
|
if shouldRetry(errExecAlreadyRunning) {
|
|
t.Fatalf("运行中标记不应重试")
|
|
}
|
|
if !shouldRetry(errors.New("模型调用失败")) {
|
|
t.Fatalf("程序报错应重试")
|
|
}
|
|
}
|