Cyrene_v2/internal/runtime/multimodal/vision/router_test.go

package vision

import (
	"context"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
)

func TestRouter_Route(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name   string
		router *Router
		target ModelInfo
		req    RouteRequest
		want   DecisionKind
		tool   string
	}{
		{
			name:   "sighted confident main -> main",
			router: NewRouter(WithVisionTool("v"), WithOCRTool("o")),
			target: ModelInfo{Name: "m", SupportsVision: true, SelfAssessment: AssessmentGood},
			want:   KindMain,
		},
		{
			name:   "poor self-assessment sighted main -> vision tool",
			router: NewRouter(WithVisionTool("v"), WithOCRTool("o")),
			target: ModelInfo{Name: "m", SupportsVision: true, SelfAssessment: AssessmentPoor},
			req:    RouteRequest{MainSelfAssessedPoor: true},
			want:   KindVisionTool,
			tool:   "v",
		},
		{
			name:   "poor self-assessment only OCR -> OCR tool",
			router: NewRouter(WithOCRTool("o")),
			target: ModelInfo{Name: "m", SupportsVision: true, SelfAssessment: AssessmentPoor},
			want:   KindOCRTool,
			tool:   "o",
		},
		{
			name:   "blind main with vision tool -> vision tool",
			router: NewRouter(WithVisionTool("v")),
			target: ModelInfo{Name: "m", SupportsVision: false},
			want:   KindVisionTool,
			tool:   "v",
		},
		{
			name:   "blind main text primary -> OCR tool",
			router: NewRouter(WithVisionTool("v"), WithOCRTool("o")),
			target: ModelInfo{Name: "m", SupportsVision: false},
			req:    RouteRequest{TextPrimary: true},
			want:   KindOCRTool,
			tool:   "o",
		},
		{
			name:   "blind main only OCR non-text -> OCR fallback",
			router: NewRouter(WithOCRTool("o")),
			target: ModelInfo{Name: "m", SupportsVision: false},
			want:   KindOCRTool,
			tool:   "o",
		},
		{
			name:   "blind main no tools -> unsupported",
			router: NewRouter(),
			target: ModelInfo{Name: "m", SupportsVision: false},
			want:   KindUnsupported,
		},
		{
			name:   "sighted main no tools -> main",
			router: NewRouter(),
			target: ModelInfo{Name: "m", SupportsVision: true, SelfAssessment: AssessmentGood},
			want:   KindMain,
		},
		{
			name:   "sighted main poor no tools -> main fallback",
			router: NewRouter(),
			target: ModelInfo{Name: "m", SupportsVision: true, SelfAssessment: AssessmentPoor},
			want:   KindMain,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got := tc.router.Route(tc.target, tc.req)
			if got.Kind != tc.want {
				t.Fatalf("Kind = %q, want %q (reason: %s)", got.Kind, tc.want, got.Reason)
			}
			if tc.tool != "" && got.Tool != tc.tool {
				t.Errorf("Tool = %q, want %q", got.Tool, tc.tool)
			}
			if got.Reason == "" {
				t.Error("Reason must be non-empty")
			}
		})
	}
}

// fakeVision 为工具测试实现 VisionAdapter。
type fakeVision struct{ desc string }

func (f *fakeVision) Describe(_ context.Context, _ ImageRef, _ string) (string, error) {
	return f.desc, nil
}

// fakeOCR 为工具测试实现 OCRAdapter。
type fakeOCR struct{ text string }

func (f *fakeOCR) Extract(_ context.Context, _ ImageRef) (string, error) {
	return f.text, nil
}

func TestRegisterTools(t *testing.T) {
	t.Parallel()
	reg := registry.New()
	v := NewVisionTool("vision_understand", &fakeVision{desc: "a cat"})
	o := NewOCRTool("ocr_extract", &fakeOCR{text: "HELLO"})

	if err := RegisterTools(reg, v, o); err != nil {
		t.Fatalf("RegisterTools error: %v", err)
	}
	if capv, err := reg.Get("vision_understand"); err != nil || capv.Kind != registry.KindTool {
		t.Fatalf("Get(vision_understand) = %+v, %v", capv, err)
	}
	if capo, err := reg.Get("ocr_extract"); err != nil || capo.Kind != registry.KindTool {
		t.Fatalf("Get(ocr_extract) = %+v, %v", capo, err)
	}
	if len(reg.List(registry.KindTool)) != 2 {
		t.Errorf("List(tool) = %d, want 2", len(reg.List(registry.KindTool)))
	}
}

func TestRegisterTools_NilRegistrar(t *testing.T) {
	t.Parallel()
	if err := RegisterTools(nil, nil, nil); err == nil {
		t.Fatal("want error for nil registrar, got nil")
	}
}

func TestVisionTool_Execute(t *testing.T) {
	t.Parallel()
	v := NewVisionTool("vision_understand", &fakeVision{desc: "a cat"})
	res, err := v.Execute(context.Background(), map[string]any{
		"image":  "data:image/png;base64,aGVsbG8=",
		"prompt": "what is this",
	})
	if err != nil {
		t.Fatalf("Execute error: %v", err)
	}
	m, ok := res.(map[string]any)
	if !ok || m["description"] != "a cat" {
		t.Errorf("result = %v, want description=a cat", res)
	}
}

func TestVisionTool_Execute_BadArgs(t *testing.T) {
	t.Parallel()
	v := NewVisionTool("vision_understand", &fakeVision{desc: ""})
	if _, err := v.Execute(context.Background(), map[string]any{"prompt": "x"}); err == nil {
		t.Fatal("want error for missing image, got nil")
	}
	if _, err := v.Execute(context.Background(), "not-a-map"); err == nil {
		t.Fatal("want error for non-map args, got nil")
	}
}

func TestOCRTool_Execute(t *testing.T) {
	t.Parallel()
	o := NewOCRTool("ocr_extract", &fakeOCR{text: "HELLO"})
	res, err := o.Execute(context.Background(), map[string]any{"image": "aGVsbG8="})
	if err != nil {
		t.Fatalf("Execute error: %v", err)
	}
	m, ok := res.(map[string]any)
	if !ok || m["text"] != "HELLO" {
		t.Errorf("result = %v, want text=HELLO", res)
	}
}