Cyrene_v2/internal/capability/mcp/registry_test.go

package mcp

import (
	"context"
	"errors"
	"testing"
	"time"

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

func sampleTools() []ToolInfo {
	return []ToolInfo{
		{Name: "add", Description: "Adds numbers", InputSchema: map[string]any{"type": "object"}},
		{Name: "sub", Description: "Subtracts numbers", InputSchema: map[string]any{"type": "object"}},
	}
}

func TestRegistry_Register_AutoRegisters(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{tools: sampleTools()}
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))

	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	for _, name := range []string{"add", "sub"} {
		cap, err := reg.Get(name)
		if err != nil {
			t.Fatalf("Get(%q): %v", name, err)
		}
		if cap.Kind != registry.KindTool {
			t.Fatalf("Kind(%q) = %q, want %q", name, cap.Kind, registry.KindTool)
		}
	}
	if !ft.hasMethod("tools/list") {
		t.Fatalf("tools/list was not called; methods = %v", ft.calledMethods())
	}
}

func TestRegistry_Register_PropagatesConnectError(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{startErr: errors.New("boom")}
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err == nil {
		t.Fatal("Register succeeded, want connect error")
	}
}

func TestRegistry_Register_DuplicateTool(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}}
	reg := registry.New()
	// Pre-register a tool with the same name so the MCP registration collides.
	_ = reg.RegisterTool(&testTool{name: "add"})
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	err := r.Register(context.Background())
	if !errors.Is(err, registry.ErrDuplicate) {
		t.Fatalf("Register error = %v, want registry.ErrDuplicate", err)
	}
}

func TestRegistry_Register_NilRegistry(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{}
	r := NewRegistry(nil, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err == nil {
		t.Fatal("Register with nil registry succeeded, want error")
	}
}

func TestRegistry_Unregister_RemovesTools(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{tools: sampleTools()}
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	if err := r.Unregister(context.Background()); err != nil {
		t.Fatalf("Unregister: %v", err)
	}
	for _, name := range []string{"add", "sub"} {
		if _, err := reg.Get(name); !errors.Is(err, registry.ErrNotFound) {
			t.Fatalf("Get(%q) after Unregister = %v, want ErrNotFound", name, err)
		}
	}
	if ft.closes != 1 {
		t.Fatalf("closes = %d, want 1 (Unregister must close the client)", ft.closes)
	}
}

func TestRegistry_Unregister_Idempotent(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{}
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	if err := r.Unregister(context.Background()); err != nil {
		t.Fatalf("first Unregister: %v", err)
	}
	if err := r.Unregister(context.Background()); err != nil {
		t.Fatalf("second Unregister: %v", err)
	}
}

func TestRegistry_ToolListChangedRefreshes(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{tools: sampleTools()}
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	// The server's tool list changes: sub is dropped, mul appears.
	ft.setTools([]ToolInfo{
		{Name: "add", Description: "Adds numbers", InputSchema: map[string]any{"type": "object"}},
		{Name: "mul", Description: "Multiplies numbers", InputSchema: map[string]any{"type": "object"}},
	})
	ft.notifyServer(notificationToolsListChanged)

	waitFor(t, time.Second, func() bool {
		_, err := reg.Get("mul")
		return err == nil
	})
	if _, err := reg.Get("add"); err != nil {
		t.Fatalf("Get(add) after refresh = %v, want present", err)
	}
	if _, err := reg.Get("sub"); !errors.Is(err, registry.ErrNotFound) {
		t.Fatalf("Get(sub) after refresh = %v, want ErrNotFound", err)
	}
}

func TestRegistry_ToolListChangedPreservesForeignTools(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{tools: sampleTools()}
	reg := registry.New()
	_ = reg.RegisterTool(&testTool{name: "add"}) // a foreign 'add' before refresh
	r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
	if err := r.Register(context.Background()); err != nil {
		// 'add' collides with the foreign tool; Register must surface the dup.
		if !errors.Is(err, registry.ErrDuplicate) {
			t.Fatalf("Register error = %v, want registry.ErrDuplicate", err)
		}
	}
	// After the collision, sub is registered, 'add' stays foreign.
	ft.setTools([]ToolInfo{{Name: "sub", Description: "Subtracts", InputSchema: map[string]any{}}})
	ft.notifyServer(notificationToolsListChanged)
	waitFor(t, time.Second, func() bool {
		_, err := reg.Get("sub")
		return err == nil
	})
	// The foreign 'add' must be preserved even though it is not in the server list.
	if _, err := reg.Get("add"); err != nil {
		t.Fatalf("Get(add) after refresh = %v, want foreign tool preserved", err)
	}
}

// testTool is a minimal sdk.Tool used to pre-register a colliding name.
type testTool struct {
	name string
}

func (t *testTool) Name() string                                  { return t.name }
func (t *testTool) Description() string                           { return "desc " + t.name }
func (t *testTool) Parameters() (map[string]any, error)           { return map[string]any{}, nil }
func (t *testTool) Execute(_ context.Context, _ any) (any, error) { return nil, nil }