package llm import ( "context" "fmt" "testing" "time" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" "github.com/sundynix/sundynix-dispatcher/internal/harness" ) // fakeModel 是 model.ToolCallingChatModel 测试替身:fail=true 则 Generate/Stream 报错。 type fakeModel struct { name string reply string fail bool calls *int // 共享计数器,记录被调次数(验证 failover 没多调/早停) } func (f *fakeModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { if f.calls != nil { *f.calls++ } if f.fail { return nil, fmt.Errorf("%s down", f.name) } return schema.AssistantMessage(f.reply, nil), nil } func (f *fakeModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { if f.calls != nil { *f.calls++ } if f.fail { return nil, fmt.Errorf("%s down", f.name) } sr, sw := schema.Pipe[*schema.Message](1) go func() { sw.Send(schema.AssistantMessage(f.reply, nil), nil); sw.Close() }() return sr, nil } func (f *fakeModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { return f, nil } func genText(t *testing.T, ctx context.Context, m model.ToolCallingChatModel) (string, error) { t.Helper() out, err := m.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) if err != nil { return "", err } return out.Content, nil } // TestFailover_PrimaryWorks 主可用 → 直接用主,备用不被调用。 func TestFailover_PrimaryWorks(t *testing.T) { var pc, fc int m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "primary", reply: "主回答", calls: &pc}, &fakeModel{name: "fb", reply: "备回答", calls: &fc}, }, nil) ans, err := genText(t, context.Background(), m) if err != nil || ans != "主回答" { t.Fatalf("应返回主回答, got %q err=%v", ans, err) } if pc != 1 || fc != 0 { t.Fatalf("主可用时不应调备用: primary=%d fb=%d", pc, fc) } } // TestFailover_PrimaryDownUseFallback 主挂 → 切备用并成功。 func TestFailover_PrimaryDownUseFallback(t *testing.T) { var pc, fc int switched := false m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "primary", fail: true, calls: &pc}, &fakeModel{name: "fb", reply: "备回答", calls: &fc}, }, func(int, error) { switched = true }) ans, err := genText(t, context.Background(), m) if err != nil || ans != "备回答" { t.Fatalf("主挂应切备用, got %q err=%v", ans, err) } if pc != 1 || fc != 1 || !switched { t.Fatalf("主备各调一次且触发回调: primary=%d fb=%d switched=%v", pc, fc, switched) } } // TestFailover_AllDown 全挂 → 返回最后一个错误。 func TestFailover_AllDown(t *testing.T) { m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "p", fail: true}, &fakeModel{name: "fb", fail: true}, }, nil) if _, err := genText(t, context.Background(), m); err == nil { t.Fatal("全挂应返回错误") } } // TestFailover_CtxCanceledNoSwitch 调用方已取消 → 不切备用(切了也没用,且违背意图)。 func TestFailover_CtxCanceledNoSwitch(t *testing.T) { var pc, fc int ctx, cancel := context.WithCancel(context.Background()) cancel() // 立即取消 m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "p", fail: true, calls: &pc}, &fakeModel{name: "fb", reply: "备", calls: &fc}, }, nil) if _, err := genText(t, ctx, m); err == nil { t.Fatal("已取消仍应返回主错误") } if fc != 0 { t.Fatalf("取消时不应切备用: fb 被调 %d 次", fc) } } // TestFailover_StreamFailover Stream 建流失败也切备用。 func TestFailover_StreamFailover(t *testing.T) { m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "p", fail: true}, &fakeModel{name: "fb", reply: "流备"}, }, nil) sr, err := m.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) if err != nil { t.Fatalf("Stream 应切到备用: %v", err) } defer sr.Close() chunk, _ := sr.Recv() if chunk == nil || chunk.Content != "流备" { t.Fatalf("应收到备用的流, got %+v", chunk) } } // 主模型持续失败达阈值 → 熔断 → 后续请求跳过主、直连备用(省掉每次白试主的失败往返)。 func TestFailover_TrippedPrimarySkipped(t *testing.T) { pc, fc := 0, 0 p := &fakeModel{name: "p", fail: true, calls: &pc} fb := &fakeModel{name: "fb", reply: "备", calls: &fc} m := newFailoverModel([]model.ToolCallingChatModel{p, fb}, nil) // 连打 threshold+2 次:前 threshold 次主被调并失败→切备;熔断后主被跳过。 for i := 0; i < fbBreakerThreshold+2; i++ { if ans, err := genText(t, context.Background(), m); err != nil || ans != "备" { t.Fatalf("每次都应最终拿到备用回答, got %q err=%v", ans, err) } } if pc != fbBreakerThreshold { t.Fatalf("主应在失败 %d 次后被熔断跳过,实际被调 %d 次", fbBreakerThreshold, pc) } if fc != fbBreakerThreshold+2 { t.Fatalf("备用应每次都被调: got %d", fc) } } // 关键:WithTools 重包必须共享熔断状态。主已熔断 → 新链仍跳过主(否则 rewrap 清零 = 熔断失效)。 func TestFailover_WithToolsSharesBreakerState(t *testing.T) { pc, fc := 0, 0 p := &fakeModel{name: "p", fail: true, calls: &pc} fb := &fakeModel{name: "fb", reply: "备", calls: &fc} m := newFailoverModel([]model.ToolCallingChatModel{p, fb}, nil) for i := 0; i < fbBreakerThreshold; i++ { // 打到主熔断 _, _ = genText(t, context.Background(), m) } if pc != fbBreakerThreshold { t.Fatalf("主应被调 %d 次, got %d", fbBreakerThreshold, pc) } bound, err := m.WithTools(nil) if err != nil { t.Fatal(err) } before := pc if _, err := genText(t, context.Background(), bound); err != nil { t.Fatal(err) } if pc != before { t.Fatalf("重包后主不应再被调(熔断状态须共享),却多调了 %d 次", pc-before) } } // 熔断后冷却到点 → 半开探测打回主 → 主已恢复则成功切回主。 func TestFailover_RecoversAfterCooldown(t *testing.T) { pc, fc := 0, 0 p := &fakeModel{name: "p", reply: "主回答", fail: true, calls: &pc} fb := &fakeModel{name: "fb", reply: "备", calls: &fc} m := &failoverModel{ models: []model.ToolCallingChatModel{p, fb}, breakers: []*harness.CircuitBreaker{ harness.NewCircuitBreakerWith(fbBreakerThreshold, 30*time.Millisecond, 1), harness.NewCircuitBreakerWith(fbBreakerThreshold, 30*time.Millisecond, 1), }, } for i := 0; i < fbBreakerThreshold; i++ { // 打到主熔断 _, _ = genText(t, context.Background(), m) } tripped := pc p.fail = false // 主恢复健康 // 冷却前:主仍被跳过。 if _, _ = genText(t, context.Background(), m); pc != tripped { t.Fatalf("冷却前主应仍被跳过,却被调用了") } // 冷却到点:半开探测打回主 → 成功恢复。 time.Sleep(40 * time.Millisecond) ans, err := genText(t, context.Background(), m) if err != nil || ans != "主回答" { t.Fatalf("冷却后应探测回主并成功, got %q err=%v", ans, err) } if pc <= tripped { t.Fatalf("冷却后主应被再次探测调用") } } // TestFailover_WithToolsBindsAll WithTools 给链上每个模型绑定,返回仍是 failover 链。 func TestFailover_WithToolsBindsAll(t *testing.T) { m := newFailoverModel([]model.ToolCallingChatModel{ &fakeModel{name: "p", fail: true}, &fakeModel{name: "fb", reply: "带工具的备"}, }, nil) bound, err := m.WithTools(nil) if err != nil { t.Fatal(err) } ans, err := genText(t, context.Background(), bound) if err != nil || ans != "带工具的备" { t.Fatalf("绑定工具后 failover 仍生效, got %q err=%v", ans, err) } }