119 lines
3.8 KiB
Go
119 lines
3.8 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
|
|
"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)
|
|
}
|
|
}
|
|
|
|
// 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 }
|