flipside-mcp-extension/main_test.go

package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
)

func TestNewMCPProxy(t *testing.T) {
	tests := []struct {
		name      string
		remoteURL string
		apiKey    string
		debug     bool
		wantURL   string
	}{
		{
			name:      "basic URL without modification",
			remoteURL: "https://example.com/mcp",
			apiKey:    "test-key",
			debug:     false,
			wantURL:   "https://example.com/mcp",
		},
		{
			name:      "URL with SSE conversion",
			remoteURL: "https://example.com/sse",
			apiKey:    "test-key",
			debug:     true,
			wantURL:   "https://example.com/mcp",
		},
		{
			name:      "URL with SSE in path conversion",
			remoteURL: "https://mcp.flipsidecrypto.xyz/beta/sse",
			apiKey:    "test-key",
			debug:     false,
			wantURL:   "https://mcp.flipsidecrypto.xyz/beta/mcp",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			proxy := NewMCPProxy(tt.remoteURL, tt.apiKey, tt.debug)

			if proxy.remoteURL != tt.wantURL {
				t.Errorf("NewMCPProxy() remoteURL = %v, want %v", proxy.remoteURL, tt.wantURL)
			}

			if proxy.apiKey != tt.apiKey {
				t.Errorf("NewMCPProxy() apiKey = %v, want %v", proxy.apiKey, tt.apiKey)
			}

			if proxy.debug != tt.debug {
				t.Errorf("NewMCPProxy() debug = %v, want %v", proxy.debug, tt.debug)
			}

			if proxy.logger == nil {
				t.Error("NewMCPProxy() logger should not be nil")
			}
		})
	}
}

func TestMCPProxy_createRemoteClient(t *testing.T) {
	tests := []struct {
		name      string
		remoteURL string
		apiKey    string
		wantErr   bool
	}{
		{
			name:      "valid URL",
			remoteURL: "https://example.com/mcp",
			apiKey:    "test-key",
			wantErr:   false,
		},
		{
			name:      "invalid URL",
			remoteURL: "://invalid-url",
			apiKey:    "test-key",
			wantErr:   true,
		},
		{
			name:      "empty API key",
			remoteURL: "https://example.com/mcp",
			apiKey:    "",
			wantErr:   false, // Empty API key is allowed, just won't authenticate
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			proxy := NewMCPProxy(tt.remoteURL, tt.apiKey, false)
			err := proxy.createRemoteClient()

			if (err != nil) != tt.wantErr {
				t.Errorf("createRemoteClient() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if !tt.wantErr && proxy.client == nil {
				t.Error("createRemoteClient() should set client when successful")
			}
		})
	}
}

func TestMaskAPIKey(t *testing.T) {
	tests := []struct {
		name   string
		apiKey string
		want   string
	}{
		{
			name:   "short key",
			apiKey: "short",
			want:   "***",
		},
		{
			name:   "normal key",
			apiKey: "abcd1234567890xyz",
			want:   "abcd...0xyz",
		},
		{
			name:   "exactly 8 chars",
			apiKey: "12345678",
			want:   "***",
		},
		{
			name:   "9 chars",
			apiKey: "123456789",
			want:   "1234...6789",
		},
		{
			name:   "empty key",
			apiKey: "",
			want:   "***",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := maskAPIKey(tt.apiKey)
			if got != tt.want {
				t.Errorf("maskAPIKey() = %v, want %v", got, tt.want)
			}
		})
	}
}

// Mock server for testing proxy functionality
func createMockMCPServer(t *testing.T) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Check authentication
		if r.Header.Get("Authorization") == "" && r.URL.Query().Get("apiKey") == "" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		// Simulate different MCP endpoints based on the request
		switch {
		case strings.Contains(r.URL.Path, "/mcp") && r.Method == "POST":
			// Read the request body to determine the MCP method
			var body map[string]interface{}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				w.WriteHeader(http.StatusBadRequest)
				return
			}

			method, ok := body["method"].(string)
			if !ok {
				// This might be a notification - just return OK
				w.WriteHeader(http.StatusOK)
				return
			}

			w.Header().Set("Content-Type", "application/json")

			switch method {
			case "initialize":
				response := map[string]interface{}{
					"jsonrpc": "2.0",
					"id":      body["id"],
					"result": map[string]interface{}{
						"protocolVersion": "2024-11-05",
						"capabilities": map[string]interface{}{
							"tools": map[string]interface{}{},
						},
						"serverInfo": map[string]interface{}{
							"name":    "mock-server",
							"version": "1.0.0",
						},
					},
				}
				if err := json.NewEncoder(w).Encode(response); err != nil {
					http.Error(w, "Internal server error", http.StatusInternalServerError)
					return
				}
			case "initialized":
				// Handle initialized notification
				w.WriteHeader(http.StatusOK)
				return

			case "tools/list":
				response := map[string]interface{}{
					"jsonrpc": "2.0",
					"id":      body["id"],
					"result": map[string]interface{}{
						"tools": []interface{}{
							map[string]interface{}{
								"name":        "test_tool",
								"description": "A test tool for unit testing",
								"inputSchema": map[string]interface{}{
									"type": "object",
									"properties": map[string]interface{}{
										"query": map[string]interface{}{
											"type":        "string",
											"description": "Test query parameter",
										},
									},
									"required": []string{"query"},
								},
							},
						},
					},
				}
				if err := json.NewEncoder(w).Encode(response); err != nil {
					http.Error(w, "Internal server error", http.StatusInternalServerError)
					return
				}

			case "tools/call":
				params, ok := body["params"].(map[string]interface{})
				if !ok {
					w.WriteHeader(http.StatusBadRequest)
					return
				}

				toolName, ok := params["name"].(string)
				if !ok || toolName != "test_tool" {
					w.WriteHeader(http.StatusBadRequest)
					return
				}

				response := map[string]interface{}{
					"jsonrpc": "2.0",
					"id":      body["id"],
					"result": map[string]interface{}{
						"content": []interface{}{
							map[string]interface{}{
								"type": "text",
								"text": "Mock tool response",
							},
						},
						"isError": false,
					},
				}
				if err := json.NewEncoder(w).Encode(response); err != nil {
					http.Error(w, "Internal server error", http.StatusInternalServerError)
					return
				}

			case "resources/list":
				response := map[string]interface{}{
					"jsonrpc": "2.0",
					"id":      body["id"],
					"result": map[string]interface{}{
						"resources": []interface{}{},
					},
				}
				if err := json.NewEncoder(w).Encode(response); err != nil {
					http.Error(w, "Internal server error", http.StatusInternalServerError)
					return
				}

			case "prompts/list":
				response := map[string]interface{}{
					"jsonrpc": "2.0",
					"id":      body["id"],
					"result": map[string]interface{}{
						"prompts": []interface{}{},
					},
				}
				if err := json.NewEncoder(w).Encode(response); err != nil {
					http.Error(w, "Internal server error", http.StatusInternalServerError)
					return
				}

			default:
				w.WriteHeader(http.StatusBadRequest)
			}
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
}

func TestMCPProxy_Integration(t *testing.T) {
	// Create a mock server
	mockServer := createMockMCPServer(t)
	defer mockServer.Close()

	// Create proxy with mock server URL
	proxy := NewMCPProxy(mockServer.URL+"/mcp", "test-api-key", true)

	// Test client creation
	err := proxy.createRemoteClient()
	if err != nil {
		t.Fatalf("Failed to create remote client: %v", err)
	}

	if proxy.client == nil {
		t.Fatal("Client should not be nil after creation")
	}
}

func TestMCPProxy_Authentication(t *testing.T) {
	// Test that API key is properly added to requests
	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Check that API key is present in header or query param
		hasAuthHeader := strings.HasPrefix(r.Header.Get("Authorization"), "Bearer test-api-key")
		hasQueryParam := r.URL.Query().Get("apiKey") == "test-api-key"

		if !hasAuthHeader && !hasQueryParam {
			t.Errorf("Expected API key in Authorization header or apiKey query param")
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		// Check User-Agent header
		if r.Header.Get("User-Agent") != "flipside-mcp-proxy/1.0" {
			t.Errorf("Expected User-Agent header to be flipside-mcp-proxy/1.0, got %s", r.Header.Get("User-Agent"))
		}

		w.WriteHeader(http.StatusOK)
		if _, err := w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)); err != nil {
			t.Errorf("Failed to write response: %v", err)
		}
	}))
	defer mockServer.Close()

	proxy := NewMCPProxy(mockServer.URL, "test-api-key", false)
	err := proxy.createRemoteClient()

	if err != nil {
		t.Fatalf("Failed to create remote client: %v", err)
	}
}

func TestMCPProxy_URLConversion(t *testing.T) {
	tests := []struct {
		name     string
		inputURL string
		wantURL  string
	}{
		{
			name:     "convert SSE to MCP",
			inputURL: "https://example.com/sse",
			wantURL:  "https://example.com/mcp",
		},
		{
			name:     "convert SSE in path to MCP",
			inputURL: "https://mcp.flipsidecrypto.xyz/beta/sse",
			wantURL:  "https://mcp.flipsidecrypto.xyz/beta/mcp",
		},
		{
			name:     "no conversion needed",
			inputURL: "https://example.com/mcp",
			wantURL:  "https://example.com/mcp",
		},
		{
			name:     "no conversion for different endpoints",
			inputURL: "https://example.com/api/v1/endpoint",
			wantURL:  "https://example.com/api/v1/endpoint",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			proxy := NewMCPProxy(tt.inputURL, "test-key", false)
			if proxy.remoteURL != tt.wantURL {
				t.Errorf("URL conversion failed: got %v, want %v", proxy.remoteURL, tt.wantURL)
			}
		})
	}
}

func TestMCPProxy_ErrorHandling(t *testing.T) {
	// Test error handling for invalid URLs
	t.Run("invalid URL", func(t *testing.T) {
		proxy := NewMCPProxy("://invalid", "test-key", false)
		err := proxy.createRemoteClient()
		if err == nil {
			t.Error("Expected error for invalid URL")
		}
		if !strings.Contains(err.Error(), "failed to parse remote URL") {
			t.Errorf("Expected URL parsing error, got: %v", err)
		}
	})

	// Test error handling for network failures
	t.Run("network failure", func(t *testing.T) {
		proxy := NewMCPProxy("https://nonexistent.example.com/mcp", "test-key", false)
		err := proxy.createRemoteClient()
		// Client creation should succeed, but connection will fail later
		if err != nil {
			t.Errorf("Unexpected error during client creation: %v", err)
		}

		// Test that the proxy handles connection failures gracefully
		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
		defer cancel()

		if proxy.client != nil {
			// This should fail due to nonexistent server
			_, err = proxy.client.Initialize(ctx, mcp.InitializeRequest{
				Params: mcp.InitializeParams{
					ProtocolVersion: "2024-11-05",
					Capabilities:    mcp.ClientCapabilities{},
					ClientInfo: mcp.Implementation{
						Name:    "test-client",
						Version: "1.0.0",
					},
				},
			})
			if err == nil {
				t.Error("Expected error when connecting to nonexistent server")
			}
		}
	})
}

func TestMCPProxy_SetupProxyServerCreation(t *testing.T) {
	// Test that the proxy can be created and client initialized
	// without requiring full MCP protocol compliance
	proxy := NewMCPProxy("https://example.com/mcp", "test-api-key", false)

	// Test client creation
	err := proxy.createRemoteClient()
	if err != nil {
		t.Fatalf("Failed to create remote client: %v", err)
	}

	if proxy.client == nil {
		t.Error("Client should not be nil after creation")
	}

	// Test server creation basics (without full initialization)
	if proxy.server != nil {
		t.Error("Server should be nil before setup")
	}
}