labmcp/internal/mcp/server_test.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"strings"
	"testing"

	"git.t-juice.club/torjus/labmcp/internal/database"
)

func TestServerInitialize(t *testing.T) {
	store := setupTestStore(t)
	server := NewServer(store, nil)

	input := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`

	resp := runRequest(t, server, input)

	if resp.Error != nil {
		t.Fatalf("Unexpected error: %v", resp.Error)
	}

	result, ok := resp.Result.(map[string]interface{})
	if !ok {
		t.Fatalf("Expected map result, got %T", resp.Result)
	}

	if result["protocolVersion"] != ProtocolVersion {
		t.Errorf("protocolVersion = %v, want %v", result["protocolVersion"], ProtocolVersion)
	}

	serverInfo := result["serverInfo"].(map[string]interface{})
	if serverInfo["name"] != "nixos-options" {
		t.Errorf("serverInfo.name = %v, want nixos-options", serverInfo["name"])
	}
}

func TestServerToolsList(t *testing.T) {
	store := setupTestStore(t)
	server := NewServer(store, nil)

	input := `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`

	resp := runRequest(t, server, input)

	if resp.Error != nil {
		t.Fatalf("Unexpected error: %v", resp.Error)
	}

	result, ok := resp.Result.(map[string]interface{})
	if !ok {
		t.Fatalf("Expected map result, got %T", resp.Result)
	}

	tools, ok := result["tools"].([]interface{})
	if !ok {
		t.Fatalf("Expected tools array, got %T", result["tools"])
	}

	// Should have 6 tools
	if len(tools) != 6 {
		t.Errorf("Expected 6 tools, got %d", len(tools))
	}

	// Check tool names
	expectedTools := map[string]bool{
		"search_options":  false,
		"get_option":      false,
		"get_file":        false,
		"index_revision":  false,
		"list_revisions":  false,
		"delete_revision": false,
	}

	for _, tool := range tools {
		toolMap := tool.(map[string]interface{})
		name := toolMap["name"].(string)
		if _, ok := expectedTools[name]; ok {
			expectedTools[name] = true
		}
	}

	for name, found := range expectedTools {
		if !found {
			t.Errorf("Tool %q not found in tools list", name)
		}
	}
}

func TestServerMethodNotFound(t *testing.T) {
	store := setupTestStore(t)
	server := NewServer(store, nil)

	input := `{"jsonrpc":"2.0","id":1,"method":"unknown/method"}`

	resp := runRequest(t, server, input)

	if resp.Error == nil {
		t.Fatal("Expected error for unknown method")
	}

	if resp.Error.Code != MethodNotFound {
		t.Errorf("Error code = %d, want %d", resp.Error.Code, MethodNotFound)
	}
}

func TestServerParseError(t *testing.T) {
	store := setupTestStore(t)
	server := NewServer(store, nil)

	input := `not valid json`

	resp := runRequest(t, server, input)

	if resp.Error == nil {
		t.Fatal("Expected parse error")
	}

	if resp.Error.Code != ParseError {
		t.Errorf("Error code = %d, want %d", resp.Error.Code, ParseError)
	}
}

func TestServerNotification(t *testing.T) {
	store := setupTestStore(t)
	server := NewServer(store, nil)

	// Notification (no response expected)
	input := `{"jsonrpc":"2.0","method":"notifications/initialized"}`

	var output bytes.Buffer
	ctx := context.Background()
	err := server.Run(ctx, strings.NewReader(input+"\n"), &output)
	if err != nil && err != io.EOF {
		t.Fatalf("Run failed: %v", err)
	}

	// Should not produce any output for notifications
	if output.Len() > 0 {
		t.Errorf("Expected no output for notification, got: %s", output.String())
	}
}

// Helper functions

func setupTestStore(t *testing.T) database.Store {
	t.Helper()

	store, err := database.NewSQLiteStore(":memory:")
	if err != nil {
		t.Fatalf("Failed to create store: %v", err)
	}

	if err := store.Initialize(context.Background()); err != nil {
		t.Fatalf("Failed to initialize store: %v", err)
	}

	t.Cleanup(func() {
		store.Close()
	})

	return store
}

func runRequest(t *testing.T, server *Server, input string) *Response {
	t.Helper()

	var output bytes.Buffer
	ctx := context.Background()

	// Run with input terminated by newline
	err := server.Run(ctx, strings.NewReader(input+"\n"), &output)
	if err != nil && err != io.EOF {
		t.Fatalf("Run failed: %v", err)
	}

	if output.Len() == 0 {
		t.Fatal("Expected response, got empty output")
	}

	var resp Response
	if err := json.Unmarshal(output.Bytes(), &resp); err != nil {
		t.Fatalf("Failed to parse response: %v\nOutput: %s", err, output.String())
	}

	return &resp
}