sundynix-agentix/sundynix-dispatcher/internal/llm/failover_test.go

package llm

import (
	"context"
	"fmt"
	"testing"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"
)

// 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)
	}
}

// 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)
	}
}