Cyrene_v2/internal/runtime/loop/mocks_test.go

package loop

import (
	"context"
	"io"
	"log/slog"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// mgResp 是一条脚本化的适配器结果：一个响应加可选错误。
type mgResp struct {
	resp model.Response
	err  error
}

// mockAdapter 是供循环测试使用的可配置 model.Adapter。它维护两份
// 独立的脚本与调用计数（一份给 Chat，一份给 ChatStream），使
// 单个实例同时服务于非流式与流式路径。它收到的每个
// 请求都被记录，使测试能断言喂给模型的
// 消息（如注入的记忆或回灌的工具结果）。
type mockAdapter struct {
	name string

	chat      []mgResp
	chatCalls int
	reqs      []model.Request

	stream      []mgResp
	chunks      [][]string
	streamCalls int
	streamReqs  []model.Request
}

func (m *mockAdapter) Name() string {
	if m.name == "" {
		return "mock"
	}
	return m.name
}

func (m *mockAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	m.reqs = append(m.reqs, req)
	if m.chatCalls >= len(m.chat) {
		return model.Response{Content: "done"}, nil
	}
	r := m.chat[m.chatCalls]
	m.chatCalls++
	return r.resp, r.err
}

func (m *mockAdapter) ChatStream(_ context.Context, req model.Request, on func(model.StreamEvent)) (model.Response, error) {
	m.streamReqs = append(m.streamReqs, req)
	i := m.streamCalls
	m.streamCalls++
	if i < len(m.chunks) {
		for _, c := range m.chunks[i] {
			if on != nil {
				on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: c})
			}
		}
	}
	if i >= len(m.stream) {
		return model.Response{Content: "done"}, nil
	}
	return m.stream[i].resp, m.stream[i].err
}

// mockGuard 是供循环测试使用的可配置 guardrail.Guardrail。
type mockGuard struct {
	inputVerdict  *guardrail.Verdict
	outputVerdict *guardrail.Verdict
	inputErr      error
	outputErr     error
}

func (m *mockGuard) CheckInput(_ context.Context, _ *model.Request) (*guardrail.Verdict, error) {
	return m.inputVerdict, m.inputErr
}

func (m *mockGuard) CheckOutput(_ context.Context, _ *model.Response) (*guardrail.Verdict, error) {
	return m.outputVerdict, m.outputErr
}

// mockTool 是供循环测试使用的最小 sdk.Tool。它支持延迟，使测试
// 能演练每次调用的超时；也支持错误，使测试能验证
// 失败的调用被回灌给模型。
type mockTool struct {
	name   string
	desc   string
	params map[string]any
	result any
	err    error
	delay  time.Duration
}

func (m *mockTool) Name() string { return m.name }

func (m *mockTool) Description() string {
	if m.desc == "" {
		return "mock tool " + m.name
	}
	return m.desc
}

func (m *mockTool) Parameters() (map[string]any, error) { return m.params, nil }

func (m *mockTool) Execute(ctx context.Context, _ any) (any, error) {
	if m.delay > 0 {
		select {
		case <-time.After(m.delay):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	return m.result, m.err
}

// testLogger 返回一个不写任何地方的 logger，保持测试输出干净。
func testLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// newTestAgent 用给定工具、无护栏地在 ad 之上构建一个 Agent。
func newTestAgent(t *testing.T, ad model.Adapter, tools *tool.Dispatcher) *agent.Agent {
	t.Helper()
	a, err := agent.New("You are a helpful assistant.", tools, nil, agent.WithModel(ad, nil))
	if err != nil {
		t.Fatalf("agent.New: %v", err)
	}
	return a
}

// dispatcherWith 用默认工具超时构建一个 Dispatcher。
func dispatcherWith(t *testing.T, tools ...sdk.Tool) *tool.Dispatcher {
	t.Helper()
	return dispatcherWithTimeout(t, 0, tools...)
}

// dispatcherWithTimeout 构建一个应用自定义默认超时的 Dispatcher
// （非正值选用包默认值）并注册工具。
func dispatcherWithTimeout(t *testing.T, timeout time.Duration, tools ...sdk.Tool) *tool.Dispatcher {
	t.Helper()
	reg := registry.New()
	var opts []tool.Option
	if timeout > 0 {
		opts = append(opts, tool.WithDefaultTimeout(timeout))
	}
	d, err := tool.New(reg, opts...)
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	for _, tk := range tools {
		if err := d.Register(tk); err != nil {
			t.Fatalf("Register %q: %v", tk.Name(), err)
		}
	}
	return d
}

// newTestMemoryStore 打开一个隔离的内存 SQLite 数据库并返回
// 基于它的 memory.Store。每次调用拥有私有 DB，因此并行测试
// 天然隔离。
func newTestMemoryStore(t *testing.T) *memory.Store {
	t.Helper()
	ctx := context.Background()
	d, err := db.Open(ctx, ":memory:", testLogger())
	if err != nil {
		t.Fatalf("db.Open: %v", err)
	}
	t.Cleanup(func() { _ = d.Close() })
	s, err := memory.NewStore(ctx, d, testLogger())
	if err != nil {
		t.Fatalf("memory.NewStore: %v", err)
	}
	return s
}

// newTestCheckpointStore 打开一个隔离的内存 SQLite 数据库并返回
// 基于它的 persist.Store。
func newTestCheckpointStore(t *testing.T) *persist.Store {
	t.Helper()
	ctx := context.Background()
	d, err := db.Open(ctx, ":memory:", testLogger())
	if err != nil {
		t.Fatalf("db.Open: %v", err)
	}
	t.Cleanup(func() { _ = d.Close() })
	s, err := persist.New(ctx, d, testLogger())
	if err != nil {
		t.Fatalf("persist.New: %v", err)
	}
	return s
}