package mcptests import ( "testing" "time" "github.com/maximhq/bifrost/core/schemas" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // ============================================================================= // ADD CLIENT TESTS // ============================================================================= func TestAddClientDuplicate(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } manager := setupMCPManager(t) // Add client clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) err := manager.AddClient(&clientConfig) require.NoError(t, err, "should add client first time") // Try to add same client again err = manager.AddClient(&clientConfig) // Should either return error or be idempotent if err == nil { clients := manager.GetClients() // Should still have reasonable number of clients (not double-added) assert.LessOrEqual(t, len(clients), 2, "should not duplicate clients excessively") } } // ============================================================================= // REMOVE CLIENT TESTS // ============================================================================= func TestRemoveClient(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) // Verify client exists clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Remove client err := manager.RemoveClient(clientID) require.NoError(t, err, "should remove client") // Verify client was removed clients = manager.GetClients() assert.Len(t, clients, 0, "should have no clients") } func TestRemoveClientInvalidID(t *testing.T) { t.Parallel() manager := setupMCPManager(t) // Try to remove non-existent client err := manager.RemoveClient("non-existent-id") assert.Error(t, err, "should error when removing non-existent client") } func TestRemoveClientMultiple(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } // Add multiple clients httpConfig1 := GetSampleHTTPClientConfig(config.HTTPServerURL) httpConfig1.ID = "client-1" httpConfig1.Name = "TestHTTPServer1" httpConfig2 := GetSampleHTTPClientConfig(config.HTTPServerURL) httpConfig2.ID = "client-2" httpConfig2.Name = "TestHTTPServer2" manager := setupMCPManager(t, httpConfig1, httpConfig2) // Verify both clients exist clients := manager.GetClients() assert.GreaterOrEqual(t, len(clients), 2, "should have at least two clients") // Remove first client err := manager.RemoveClient("client-1") require.NoError(t, err, "should remove first client") // Verify only one client remains clients = manager.GetClients() assert.Len(t, clients, 1, "should have one client remaining") // Remove second client err = manager.RemoveClient("client-2") require.NoError(t, err, "should remove second client") // Verify no clients remain clients = manager.GetClients() assert.Len(t, clients, 0, "should have no clients") } func TestRemoveClientDuringExecution(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) bifrost := setupBifrost(t) bifrost.SetMCPManager(manager) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Start a tool execution (if delay tool is available) ctx := createTestContext() toolCall := GetSampleEchoToolCall("call-1", "test") // Execute tool asynchronously done := make(chan bool) go func() { _, _ = bifrost.ExecuteChatMCPTool(ctx, &toolCall) done <- true }() // Small delay to let execution start time.Sleep(100 * time.Millisecond) // Remove client during execution err := manager.RemoveClient(clientID) require.NoError(t, err, "should remove client even during execution") // Wait for execution to complete <-done // Verify client was removed clients = manager.GetClients() assert.Len(t, clients, 0, "client should be removed") } // ============================================================================= // EDIT CLIENT TESTS // ============================================================================= func TestEditClient(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Edit client configuration updatedConfig := clientConfig updatedConfig.Name = "UpdatedName" updatedConfig.ToolsToExecute = []string{"calculator", "echo"} err := manager.UpdateClient(clientID, &updatedConfig) require.NoError(t, err, "should edit client") // Verify changes clients = manager.GetClients() require.Len(t, clients, 1, "should still have one client") assert.Equal(t, "UpdatedName", clients[0].Name) } func TestEditClientInvalidID(t *testing.T) { t.Parallel() manager := setupMCPManager(t) // Try to edit non-existent client clientConfig := GetSampleHTTPClientConfig("http://example.com") err := manager.UpdateClient("non-existent-id", &clientConfig) assert.Error(t, err, "should error when editing non-existent client") } func TestEditClientInvalidConfig(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Try to edit with invalid config (missing ConnectionString) invalidConfig := schemas.MCPClientConfig{ ID: clientConfig.ID, ConnectionType: schemas.MCPConnectionTypeHTTP, // Missing ConnectionString } err := manager.UpdateClient(clientID, &invalidConfig) // Should return error or leave client unchanged if err == nil { clients = manager.GetClients() if len(clients) > 0 { // Client might be in error state t.Log("Edit with invalid config did not error, checking client state") } } else { assert.Error(t, err, "should error with invalid config") } } func TestEditClientChangeConnectionType(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Try to change connection type updatedConfig := clientConfig updatedConfig.ConnectionType = schemas.MCPConnectionTypeSSE err := manager.UpdateClient(clientID, &updatedConfig) assert.Error(t, err, "should not allow connection type change") clients = manager.GetClients() if len(clients) > 0 { assert.Equal(t, schemas.MCPConnectionTypeHTTP, clients[0].ConnectionInfo.Type) } } // ============================================================================= // GET CLIENTS TESTS // ============================================================================= func TestGetMCPClients(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) // Get clients clients := manager.GetClients() assert.NotNil(t, clients, "clients should not be nil") assert.Len(t, clients, 1, "should have one client") } func TestGetMCPClientsEmpty(t *testing.T) { t.Parallel() manager := setupMCPManager(t) // Get clients when none exist clients := manager.GetClients() assert.NotNil(t, clients, "clients should not be nil") assert.Len(t, clients, 0, "should have no clients") } func TestGetMCPClientsMultiple(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } // Create HTTP client httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) httpConfig.ID = "http-get-test" applyTestConfigHeaders(t, &httpConfig) manager := setupMCPManager(t, httpConfig) // Register a tool to create the InProcess client automatically testToolHandler := func(args any) (string, error) { return "test response", nil } testTool := GetSampleEchoTool() testTool.Function.Name = "test_tool" err := manager.RegisterTool("test_tool", "Test tool", testToolHandler, testTool) require.NoError(t, err, "should register tool") // Get all clients - should have HTTP + InProcess (auto-created) clients := manager.GetClients() assert.GreaterOrEqual(t, len(clients), 2, "should have HTTP and InProcess clients") // Verify client types hasHTTP := false hasInProcess := false for _, client := range clients { if client.ConnectionInfo.Type == schemas.MCPConnectionTypeHTTP { hasHTTP = true } if client.ConnectionInfo.Type == schemas.MCPConnectionTypeInProcess { hasInProcess = true } } assert.True(t, hasHTTP, "should have HTTP client") assert.True(t, hasInProcess, "should have InProcess client (auto-created)") } // ============================================================================= // RECONNECT CLIENT TESTS // ============================================================================= func TestReconnectClient(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) applyTestConfigHeaders(t, &clientConfig) manager := setupMCPManager(t, clientConfig) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Reconnect client err := manager.ReconnectClient(clientID) require.NoError(t, err, "should reconnect client") // Verify client is still connected time.Sleep(time.Second) clients = manager.GetClients() AssertClientState(t, clients, clientID, schemas.MCPConnectionStateConnected) } func TestReconnectClientInvalidID(t *testing.T) { t.Parallel() manager := setupMCPManager(t) // Try to reconnect non-existent client err := manager.ReconnectClient("non-existent-id") assert.Error(t, err, "should error when reconnecting non-existent client") } func TestReconnectClientAfterRemoval(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) manager := setupMCPManager(t, clientConfig) clients := manager.GetClients() require.Len(t, clients, 1, "should have one client") clientID := clients[0].ExecutionConfig.ID // Remove client err := manager.RemoveClient(clientID) require.NoError(t, err, "should remove client") // Try to reconnect removed client err = manager.ReconnectClient(clientID) assert.Error(t, err, "should not reconnect removed client") } // ============================================================================= // CONCURRENT CLIENT OPERATIONS TESTS // ============================================================================= func TestConcurrentClientOperations(t *testing.T) { t.Parallel() config := GetTestConfig(t) if config.HTTPServerURL == "" { t.Skip("MCP_HTTP_URL not set") } manager := setupMCPManager(t) // Perform concurrent add operations done := make(chan bool, 5) errors := make(chan error, 5) for i := 0; i < 5; i++ { go func(id int) { clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL) clientConfig.ID = string(rune('a'+id)) + "-concurrent-client" clientConfig.Name = "TestHTTPServer" + string(rune('a'+id)) err := manager.AddClient(&clientConfig) if err != nil { errors <- err } done <- true }(i) } // Wait for all operations for i := 0; i < 5; i++ { <-done } close(errors) // Check for errors errorCount := 0 for err := range errors { t.Logf("Concurrent add error: %v", err) errorCount++ } // Most operations should succeed assert.LessOrEqual(t, errorCount, 2, "should have few errors in concurrent operations") // Verify clients were added clients := manager.GetClients() assert.GreaterOrEqual(t, len(clients), 3, "should have added multiple clients") }