Skip to content

Conversation

@ibetitsmike
Copy link

Summary

This PR creates a proper testing package to address the tech debt identified in #73. The package consolidates scattered test helpers and provides well-structured components that reduce boilerplate and improve test readability.

New Components

MockUpstreamServer

Simulates LLM provider API responses using txtar fixtures. Features:

  • Automatic streaming/non-streaming response handling
  • Response mutation for multi-turn conversations
  • Custom status code support
  • Call counting

MockMCPServer

Provides a mock MCP server with integrated call tracking. This addresses the callAccumulator issue mentioned in #73:

  • Combines server, HTTP endpoint, and call accumulator into a single struct
  • GetToolCalls(toolName) returns all recorded invocations
  • TotalCallCount() returns total invocations across all tools
  • NewMCPManager() creates an initialized ServerProxyManager ready to use

MockRecorder

Implements aibridge.Recorder for testing with convenient accessors:

  • Interceptions(), TokenUsages(), ToolUsages(), PromptUsages()
  • TotalInputTokens(), TotalOutputTokens()
  • VerifyAllInterceptionsEnded(t) for validation

TestBridge

The main orchestrator that combines everything:

  • Creates bridge, HTTP server, recorder, and optional MCP server
  • DoAnthropicRequest(t, reqBody) and DoOpenAIRequest(t, reqBody) methods
  • Automatic cleanup on test completion

Example Usage

Before (from bridge_integration_test.go):

// Manual txtar parsing
arc := txtar.Parse(fixture)
files := filesMap(arc)
reqBody := files[fixtureRequest]

// Manual server setup
ctx, cancel := context.WithTimeout(t.Context(), time.Second*30)
srv := newMockServer(ctx, t, files, nil)
t.Cleanup(srv.Close)

// Manual recorder setup
recorderClient := &mockRecorderClient{}

// Manual MCP setup with separate callAccumulator
mcpProxiers, acc := setupMCPServerProxiesForTest(t, testTracer)
mcpMgr := mcp.NewServerProxyManager(mcpProxiers, testTracer)
require.NoError(t, mcpMgr.Init(ctx))

// Manual bridge setup
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
providers := []aibridge.Provider{aibridge.NewAnthropicProvider(cfg, nil)}
b, err := aibridge.NewRequestBridge(ctx, providers, recorderClient, mcpMgr, logger, nil, testTracer)

// Manual HTTP server setup
mockSrv := httptest.NewUnstartedServer(b)
mockSrv.Config.BaseContext = func(_ net.Listener) context.Context {
    return aibridge.AsActor(ctx, userID, nil)
}
mockSrv.Start()
t.Cleanup(mockSrv.Close)

// Make request
req := createAnthropicMessagesReq(t, mockSrv.URL, reqBody)
resp, err := client.Do(req)

// Verify tool calls via separate accumulator
invocations := acc.getCallsByTool(mockToolName)

After (with aibtest):

// 1. Create mock upstream
upstream := aibtest.NewMockUpstreamServer(t, ctx, fixture)

// 2. Create MCP server with integrated call tracking
mcpServer := aibtest.NewMockMCPServer(t, aibtest.TestTracer())

// 3. Create test bridge (combines everything)
bridge := aibtest.NewTestBridge(t, ctx, aibtest.TestBridgeOptions{
    Provider:  aibridge.NewAnthropicProvider(aibtest.AnthropicConfig(upstream.URL, aibtest.DefaultAPIKey), nil),
    MCPServer: mcpServer,
})

// 4. Make request
resp := bridge.DoAnthropicRequest(t, upstream.Files()[aibtest.FixtureRequest])

// 5. Verify - tool calls are tracked in mcpServer directly
calls := mcpServer.GetToolCalls("my_tool")
bridge.Recorder.VerifyAllInterceptionsEnded(t)

Issues Addressed from #73

  • callAccumulator is now integrated into MockMCPServer
  • setupMCPServerProxiesForTest return values are consolidated
  • proxies return value complexity is eliminated
  • ✅ Common test patterns are abstracted into TestBridge
  • ✅ Package can be used externally (e.g., in coder/coder integration tests)

Testing

All existing tests continue to pass:

ok      github.com/coder/aibridge       0.186s
ok      github.com/coder/aibridge/aibtest       1.346s
ok      github.com/coder/aibridge/buildinfo     0.004s
ok      github.com/coder/aibridge/mcp   0.008s
ok      github.com/coder/aibridge/utils 0.004s

Resolves #73

This addresses the tech debt identified in issue #73 by:

1. Creating mockMCPServer struct that combines:
   - MCP server proxies (Proxies map)
   - Call tracking (previously separate callAccumulator)

2. Replacing setupMCPServerProxiesForTest with newMockMCPServer:
   - Returns single struct instead of multiple values
   - Integrates all MCP server setup into one place

3. Updating setupInjectedToolTest:
   - Returns 3 values instead of 4 (recorder, mcpMock, resp)
   - Callers use mcpMock.GetToolCalls() instead of separate accumulator

4. Adding createMockMCPSrvHandler for cases needing custom setup
   (e.g., custom server name in trace tests)

The changes reduce cognitive overhead by keeping related data together
and simplifying function signatures.

Resolves #73
@ibetitsmike ibetitsmike force-pushed the fix-aibridge-issue-73 branch from a4b35fd to 1728be9 Compare December 5, 2025 15:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Create proper testing package

1 participant