tools/registry_test.go

package tools

import (
	"context"
	"testing"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/stretchr/testify/assert"
)

// mockTool implements the Tool interface for testing
type mockTool struct {
	name        string
	description string
	schema      map[string]any
}

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

func (t *mockTool) Description() string {
	return t.description
}

func (t *mockTool) InputSchema() map[string]any {
	return t.schema
}

func (t *mockTool) Execute(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	return mcp.NewToolResultText("mock result"), nil
}

func TestNewRegistry(t *testing.T) {
	registry := NewRegistry()
	assert.NotNil(t, registry)
	assert.Equal(t, 0, registry.Count())
}

func TestRegister(t *testing.T) {
	registry := NewRegistry()

	// Test registering a valid tool
	tool := &mockTool{
		name:        "test-tool",
		description: "A test tool",
		schema: map[string]any{
			"type": "object",
		},
	}

	err := registry.Register(tool)
	assert.NoError(t, err)
	assert.Equal(t, 1, registry.Count())

	// Test registering nil tool
	err = registry.Register(nil)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "cannot register nil tool")

	// Test registering tool with empty name
	emptyNameTool := &mockTool{
		name:        "",
		description: "A tool with no name",
	}
	err = registry.Register(emptyNameTool)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "tool name cannot be empty")

	// Test registering duplicate tool
	err = registry.Register(tool)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "already registered")
}

func TestUnregister(t *testing.T) {
	registry := NewRegistry()

	tool := &mockTool{
		name:        "test-tool",
		description: "A test tool",
	}

	// Register the tool first
	err := registry.Register(tool)
	assert.NoError(t, err)

	// Test unregistering existing tool
	err = registry.Unregister("test-tool")
	assert.NoError(t, err)
	assert.Equal(t, 0, registry.Count())

	// Test unregistering non-existent tool
	err = registry.Unregister("non-existent")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not found")
}

func TestGet(t *testing.T) {
	registry := NewRegistry()

	tool := &mockTool{
		name:        "test-tool",
		description: "A test tool",
	}

	// Register the tool first
	err := registry.Register(tool)
	assert.NoError(t, err)

	// Test getting existing tool
	gotTool, exists := registry.Get("test-tool")
	assert.True(t, exists)
	assert.Equal(t, tool, gotTool)

	// Test getting non-existent tool
	_, exists = registry.Get("non-existent")
	assert.False(t, exists)
}

func TestList(t *testing.T) {
	registry := NewRegistry()

	tool1 := &mockTool{
		name:        "tool-1",
		description: "First tool",
	}
	tool2 := &mockTool{
		name:        "tool-2",
		description: "Second tool",
	}

	// Register both tools
	err := registry.Register(tool1)
	assert.NoError(t, err)
	err = registry.Register(tool2)
	assert.NoError(t, err)

	// Test listing tools
	tools := registry.List()
	assert.Len(t, tools, 2)

	// Verify tool names and descriptions
	names := make(map[string]string)
	for _, tool := range tools {
		names[tool.Name] = tool.Description
	}

	assert.Equal(t, "First tool", names["tool-1"])
	assert.Equal(t, "Second tool", names["tool-2"])
}

func TestNames(t *testing.T) {
	registry := NewRegistry()

	tool1 := &mockTool{
		name:        "tool-1",
		description: "First tool",
	}
	tool2 := &mockTool{
		name:        "tool-2",
		description: "Second tool",
	}

	// Register both tools
	err := registry.Register(tool1)
	assert.NoError(t, err)
	err = registry.Register(tool2)
	assert.NoError(t, err)

	// Test getting tool names
	names := registry.Names()
	assert.Len(t, names, 2)
	assert.Contains(t, names, "tool-1")
	assert.Contains(t, names, "tool-2")
}

func TestExecute(t *testing.T) {
	registry := NewRegistry()

	tool := &mockTool{
		name:        "test-tool",
		description: "A test tool",
	}

	// Register the tool first
	err := registry.Register(tool)
	assert.NoError(t, err)

	// Test executing existing tool
	request := mcp.CallToolRequest{}

	result, err := registry.Execute(context.Background(), "test-tool", request)
	assert.NoError(t, err)
	assert.NotNil(t, result)

	// Test executing non-existent tool
	_, err = registry.Execute(context.Background(), "non-existent", request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not found")
}

func TestClear(t *testing.T) {
	registry := NewRegistry()

	tool1 := &mockTool{
		name:        "tool-1",
		description: "First tool",
	}
	tool2 := &mockTool{
		name:        "tool-2",
		description: "Second tool",
	}

	// Register both tools
	err := registry.Register(tool1)
	assert.NoError(t, err)
	err = registry.Register(tool2)
	assert.NoError(t, err)

	// Test clearing registry
	registry.Clear()
	assert.Equal(t, 0, registry.Count())

	// Verify tools are gone
	_, exists := registry.Get("tool-1")
	assert.False(t, exists)
	_, exists = registry.Get("tool-2")
	assert.False(t, exists)
}