From f440cbd2059978123f0349d9bb42be7e4fb24ad3 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Tue, 9 Jun 2026 10:24:33 +0100 Subject: [PATCH] refactor(aibridge): move shared mock helpers to testutil (#25999) Moves the shared aibridge mock test helpers (`MockUpstream`, `MockServerProxier`, `StubToolCaller`, and the `NewFixtureResponse`/`NewFixtureToolResponse` constructors) out of `aibridge/internal/integrationtest` into `aibridge/internal/testutil`, and exports them. This lets the per-interceptor test packages (`messages`, `chatcompletions`, `responses`) reuse one set of mocks instead of each redefining its own. The symbols are exported and call sites updated, with no behavior change. Closes: https://linear.app/codercom/issue/AIGOV-397/move-mockserverproxier-into-a-shared-testutil-package > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira --- .../chatcompletions/blocking_internal_test.go | 46 +----------- .../streaming_internal_test.go | 6 +- .../intercept/messages/base_internal_test.go | 68 +++++------------ .../messages/blocking_internal_test.go | 6 +- .../messages/streaming_internal_test.go | 16 +--- .../responses/blocking_internal_test.go | 46 +----------- .../responses/streaming_internal_test.go | 6 +- .../integrationtest/apidump_internal_test.go | 2 +- .../integrationtest/bridge_internal_test.go | 72 +++++++++--------- .../integrationtest/metrics_internal_test.go | 10 +-- .../responses_internal_test.go | 10 +-- .../internal/integrationtest/setupbridge.go | 6 +- .../integrationtest/trace_internal_test.go | 10 +-- .../internal/testutil/mockserverproxier.go | 47 ++++++++++++ .../mockupstream.go | 75 ++++++++++--------- 15 files changed, 179 insertions(+), 247 deletions(-) create mode 100644 aibridge/internal/testutil/mockserverproxier.go rename aibridge/internal/{integrationtest => testutil}/mockupstream.go (81%) diff --git a/aibridge/intercept/chatcompletions/blocking_internal_test.go b/aibridge/intercept/chatcompletions/blocking_internal_test.go index 2b9afaadea..088fa4177d 100644 --- a/aibridge/intercept/chatcompletions/blocking_internal_test.go +++ b/aibridge/intercept/chatcompletions/blocking_internal_test.go @@ -1,7 +1,6 @@ package chatcompletions import ( - "context" "io" "net/http" "net/http/httptest" @@ -10,7 +9,6 @@ import ( "testing" "github.com/google/uuid" - mcplib "github.com/mark3labs/mcp-go/mcp" "github.com/openai/openai-go/v3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -450,10 +448,10 @@ func TestBlockingInterception_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's tool_use // response will reference. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", @@ -483,41 +481,3 @@ func TestBlockingInterception_AgenticLoopFailover(t *testing.T) { }) } } - -// mockServerProxier is a test implementation of mcp.ServerProxier. -type mockServerProxier struct { - tools []*mcp.Tool -} - -func (*mockServerProxier) Init(context.Context) error { - return nil -} - -func (*mockServerProxier) Shutdown(context.Context) error { - return nil -} - -func (m *mockServerProxier) ListTools() []*mcp.Tool { - return m.tools -} - -func (m *mockServerProxier) GetTool(id string) *mcp.Tool { - for _, t := range m.tools { - if t.ID == id { - return t - } - } - return nil -} - -func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) { - return nil, nil //nolint:nilnil // mock: no-op implementation -} - -// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed -// text result, so the agentic continuation can proceed. -type stubToolCaller struct{} - -func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) { - return mcplib.NewToolResultText("tool result"), nil -} diff --git a/aibridge/intercept/chatcompletions/streaming_internal_test.go b/aibridge/intercept/chatcompletions/streaming_internal_test.go index 9561c0948a..b836e7b7cc 100644 --- a/aibridge/intercept/chatcompletions/streaming_internal_test.go +++ b/aibridge/intercept/chatcompletions/streaming_internal_test.go @@ -580,10 +580,10 @@ func TestStreamingInterception_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's tool_calls // chunks will reference. The stub caller returns a // fixed text result. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index ef130deca1..ce186057c0 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -9,7 +9,6 @@ import ( "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/shared/constant" - mcpgo "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -17,6 +16,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/internal/testutil" "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/utils" @@ -476,7 +476,7 @@ func TestInjectTools_CacheBreakpoints(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tools":[`+ `{"name":"existing_tool","type":"custom","input_schema":{"type":"object","properties":{}},"cache_control":{"type":"ephemeral"}}]}`), - mcpProxy: &mockServerProxier{tools: nil}, + mcpProxy: &testutil.MockServerProxier{Tools: nil}, logger: slog.Make(), } @@ -496,8 +496,8 @@ func TestInjectTools_CacheBreakpoints(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tools":[`+ `{"name":"existing_tool","type":"custom","input_schema":{"type":"object","properties":{}},"cache_control":{"type":"ephemeral"}}]}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{ + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, }, }, @@ -525,8 +525,8 @@ func TestInjectTools_CacheBreakpoints(t *testing.T) { reqPayload: mustMessagesPayload(t, `{"tools":[`+ `{"name":"tool_with_cache_1","type":"custom","input_schema":{"type":"object","properties":{}},"cache_control":{"type":"ephemeral"}},`+ `{"name":"tool_with_cache_2","type":"custom","input_schema":{"type":"object","properties":{}}}]}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{ + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, }, }, @@ -554,8 +554,8 @@ func TestInjectTools_CacheBreakpoints(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tools":[`+ `{"name":"existing_tool_no_cache","type":"custom","input_schema":{"type":"object","properties":{}}}]}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{ + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, }, }, @@ -583,7 +583,7 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tool_choice":{"type":"auto"}}`), - mcpProxy: &mockServerProxier{tools: nil}, // No tools to inject. + mcpProxy: &testutil.MockServerProxier{Tools: nil}, // No tools to inject. logger: slog.Make(), } @@ -600,8 +600,8 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, }, logger: slog.Make(), } @@ -619,8 +619,8 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tool_choice":{"type":"auto"}}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, }, logger: slog.Make(), } @@ -638,8 +638,8 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tool_choice":{"type":"any"}}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, }, logger: slog.Make(), } @@ -657,8 +657,8 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tool_choice":{"type":"tool","name":"specific_tool"}}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, }, logger: slog.Make(), } @@ -676,8 +676,8 @@ func TestInjectTools_ParallelToolCalls(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"tool_choice":{"type":"none"}}`), - mcpProxy: &mockServerProxier{ - tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + mcpProxy: &testutil.MockServerProxier{ + Tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, }, logger: slog.Make(), } @@ -936,36 +936,6 @@ func mustMessagesPayload(t *testing.T, requestBody string) RequestPayload { return payload } -// mockServerProxier is a test implementation of mcp.ServerProxier. -type mockServerProxier struct { - tools []*mcp.Tool -} - -func (*mockServerProxier) Init(context.Context) error { - return nil -} - -func (*mockServerProxier) Shutdown(context.Context) error { - return nil -} - -func (m *mockServerProxier) ListTools() []*mcp.Tool { - return m.tools -} - -func (m *mockServerProxier) GetTool(id string) *mcp.Tool { - for _, t := range m.tools { - if t.ID == id { - return t - } - } - return nil -} - -func (*mockServerProxier) CallTool(context.Context, string, any) (*mcpgo.CallToolResult, error) { - return nil, nil //nolint:nilnil // mock: no-op implementation -} - func TestFilterBedrockBetaFlags(t *testing.T) { t.Parallel() diff --git a/aibridge/intercept/messages/blocking_internal_test.go b/aibridge/intercept/messages/blocking_internal_test.go index 9b3f0d447b..e5c3c9f6ce 100644 --- a/aibridge/intercept/messages/blocking_internal_test.go +++ b/aibridge/intercept/messages/blocking_internal_test.go @@ -443,10 +443,10 @@ func TestBlockingInterception_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's tool_use // response will reference. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", diff --git a/aibridge/intercept/messages/streaming_internal_test.go b/aibridge/intercept/messages/streaming_internal_test.go index 40b8344777..978af0db87 100644 --- a/aibridge/intercept/messages/streaming_internal_test.go +++ b/aibridge/intercept/messages/streaming_internal_test.go @@ -1,7 +1,6 @@ package messages import ( - "context" "io" "net/http" "net/http/httptest" @@ -10,7 +9,6 @@ import ( "testing" "github.com/google/uuid" - mcplib "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" @@ -367,14 +365,6 @@ data: {"type":"message_stop"} ` ) -// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed -// text result, so the agentic continuation can proceed. -type stubToolCaller struct{} - -func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) { - return mcplib.NewToolResultText("tool result"), nil -} - // TestStreamingInterception_AgenticLoopFailover covers the // scenarios that span an agentic-loop continuation: the initial // client request and the subsequent tool-call continuation can @@ -537,10 +527,10 @@ func TestStreamingInterception_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's tool_use event // will reference. The stub caller returns a fixed // text result. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", diff --git a/aibridge/intercept/responses/blocking_internal_test.go b/aibridge/intercept/responses/blocking_internal_test.go index 94acf0deef..69e682b23c 100644 --- a/aibridge/intercept/responses/blocking_internal_test.go +++ b/aibridge/intercept/responses/blocking_internal_test.go @@ -1,7 +1,6 @@ package responses import ( - "context" "io" "net/http" "net/http/httptest" @@ -10,7 +9,6 @@ import ( "testing" "github.com/google/uuid" - mcplib "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" @@ -440,10 +438,10 @@ func TestBlockingResponsesInterceptor_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's function_call // response will reference. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", @@ -473,41 +471,3 @@ func TestBlockingResponsesInterceptor_AgenticLoopFailover(t *testing.T) { }) } } - -// mockServerProxier is a test implementation of mcp.ServerProxier. -type mockServerProxier struct { - tools []*mcp.Tool -} - -func (*mockServerProxier) Init(context.Context) error { - return nil -} - -func (*mockServerProxier) Shutdown(context.Context) error { - return nil -} - -func (m *mockServerProxier) ListTools() []*mcp.Tool { - return m.tools -} - -func (m *mockServerProxier) GetTool(id string) *mcp.Tool { - for _, t := range m.tools { - if t.ID == id { - return t - } - } - return nil -} - -func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) { - return nil, nil //nolint:nilnil // mock: no-op implementation -} - -// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed -// text result, so the agentic continuation can proceed. -type stubToolCaller struct{} - -func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) { - return mcplib.NewToolResultText("tool result"), nil -} diff --git a/aibridge/intercept/responses/streaming_internal_test.go b/aibridge/intercept/responses/streaming_internal_test.go index 4f20d76c17..7c49140bfa 100644 --- a/aibridge/intercept/responses/streaming_internal_test.go +++ b/aibridge/intercept/responses/streaming_internal_test.go @@ -484,10 +484,10 @@ func TestStreamingResponsesInterceptor_AgenticLoopFailover(t *testing.T) { // Mock proxy with a tool the upstream's function_call // response will reference. The stub caller returns a // fixed text result. - proxy := &mockServerProxier{ - tools: []*mcp.Tool{ + proxy := &testutil.MockServerProxier{ + Tools: []*mcp.Tool{ { - Client: stubToolCaller{}, + Client: testutil.StubToolCaller{}, ID: "test_tool", Name: "test_tool", ServerName: "coder", diff --git a/aibridge/internal/integrationtest/apidump_internal_test.go b/aibridge/internal/integrationtest/apidump_internal_test.go index 738c245569..42811cb362 100644 --- a/aibridge/internal/integrationtest/apidump_internal_test.go +++ b/aibridge/internal/integrationtest/apidump_internal_test.go @@ -121,7 +121,7 @@ func TestAPIDump(t *testing.T) { // Setup mock upstream server. fix := fixtures.Parse(t, tc.fixture) - srv := newMockUpstream(ctx, t, newFixtureResponse(fix)) + srv := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) // Create temp dir for API dumps. dumpDir := t.TempDir() diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index dcf96134aa..ef226db2b9 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -89,7 +89,7 @@ func TestAnthropicMessages(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -139,7 +139,7 @@ func TestAnthropicMessages(t *testing.T) { assert.Equal(t, "read the foo file", promptUsages[0].Prompt) // Verify PRM attribution is NOT present on non-Bedrock Anthropic requests. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) ua := received[0].Header.Get("User-Agent") assert.NotContains(t, ua, "sdk-ua-app-id", @@ -161,9 +161,9 @@ func TestAnthropicMessages(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.AntSingleInjectedToolNoPreamble) - upstream := newMockUpstream(ctx, t, - newFixtureResponse(fix), - newErrorResponse(http.StatusInternalServerError), + upstream := testutil.NewMockUpstream(ctx, t, + testutil.NewFixtureResponse(fix), + testutil.NewErrorResponse(http.StatusInternalServerError, ""), ) mockMCP := setupMCPForTest(t, defaultTracer) @@ -274,7 +274,7 @@ func TestAnthropicMessagesModelThoughts(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -340,7 +340,7 @@ func TestAWSBedrockIntegration(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) // We define region here to validate that with Region & BaseURL defined, the latter takes precedence. bedrockCfg := &config.AWSBedrock{ @@ -373,7 +373,7 @@ func TestAWSBedrockIntegration(t *testing.T) { // Verify that Bedrock-specific model name was used in the request to the mock server // and the interception data. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) // The Anthropic SDK's Bedrock middleware extracts "model" and "stream" @@ -471,7 +471,7 @@ func TestAWSBedrockIntegration(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.AntSimpleBedrock) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bCfg := &config.AWSBedrock{ Region: "us-west-2", @@ -499,7 +499,7 @@ func TestAWSBedrockIntegration(t *testing.T) { _, err = io.ReadAll(resp.Body) require.NoError(t, err) - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) body := received[0].Body @@ -561,7 +561,7 @@ func TestAWSBedrockIntegration(t *testing.T) { // Mock Bedrock endpoint (simulates AWS). The OnRequest callback // re-signs the received request using only the declared // SignedHeaders and stores whether the signatures match. - fixResp := newFixtureResponse(fix) + fixResp := testutil.NewFixtureResponse(fix) fixResp.OnRequest = func(r *http.Request, body []byte) { authHeader := r.Header.Get("Authorization") // Passthrough requests have no SigV4 auth; skip verification. @@ -602,7 +602,7 @@ func TestAWSBedrockIntegration(t *testing.T) { recomputedSig := extractSigV4Field(verifyReq.Header.Get("Authorization"), "Signature=") signatureValid.Store(originalSig == recomputedSig) } - mockBedrock := newMockUpstream(ctx, t, fixResp) + mockBedrock := testutil.NewMockUpstream(ctx, t, fixResp) mockBedrock.AllowOverflow = true // Simulated egress proxy: modifies X-Forwarded-For and @@ -666,7 +666,7 @@ func TestAWSBedrockIntegration(t *testing.T) { defer resp.Body.Close() _, _ = io.ReadAll(resp.Body) - received := mockBedrock.receivedRequests() + received := mockBedrock.ReceivedRequests() require.NotEmpty(t, received) last := received[len(received)-1] @@ -716,7 +716,7 @@ func TestOpenAIChatCompletions(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.OaiChatSingleBuiltinTool) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -794,7 +794,7 @@ func TestOpenAIChatCompletions(t *testing.T) { // Setup mock server for multi-turn interaction. // First request → tool call response, second → tool response. fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix), newFixtureToolResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix), testutil.NewFixtureToolResponse(fix)) // Setup MCP proxies with the tool from the fixture mockMCP := setupMCPForTest(t, defaultTracer) @@ -996,7 +996,7 @@ func TestSimple(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL+tc.basePath) @@ -1009,7 +1009,7 @@ func TestSimple(t *testing.T) { require.Equal(t, http.StatusOK, resp.StatusCode) // Then: I expect the upstream request to have the correct path. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) require.Equal(t, tc.expectedPath, received[0].Path) @@ -1113,7 +1113,7 @@ func TestSessionIDTracking(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withProvider(config.ProviderAnthropic)) reqBody := fix.Request() @@ -1198,7 +1198,7 @@ func TestFallthrough(t *testing.T) { t.Parallel() fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(t.Context(), t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(t.Context(), t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL+tc.basePath) resp, err := bridgeServer.makeRequest(t, http.MethodGet, tc.requestPath, nil) @@ -1209,7 +1209,7 @@ func TestFallthrough(t *testing.T) { // Verify upstream received the request at the expected path // with the API key header. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) require.Equal(t, tc.expectedUpstreamPath, received[0].Path) require.Contains(t, received[0].Header.Get(tc.expectAuthHeader), apiKey) @@ -1538,7 +1538,7 @@ func TestErrorHandling(t *testing.T) { // Setup mock server. Error fixtures contain raw HTTP // responses that may cause the bridge to retry. fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -1610,7 +1610,7 @@ func TestErrorHandling(t *testing.T) { // Setup mock server. fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) upstream.StatusCode = http.StatusInternalServerError bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -1664,11 +1664,11 @@ func TestStableRequestEncoding(t *testing.T) { // Create a mock upstream that serves the same blocking response for each request. count := 10 - responses := make([]upstreamResponse, count) + responses := make([]testutil.UpstreamResponse, count) for i := range count { - responses[i] = newFixtureResponse(fix) + responses[i] = testutil.NewFixtureResponse(fix) } - upstream := newMockUpstream(ctx, t, responses...) + upstream := testutil.NewMockUpstream(ctx, t, responses...) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withMCP(mockMCP), @@ -1683,7 +1683,7 @@ func TestStableRequestEncoding(t *testing.T) { } // All upstream request bodies should be identical. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, count) reference := string(received[0].Body) for _, r := range received[1:] { @@ -1931,7 +1931,7 @@ func TestAnthropicToolChoiceParallelDisabled(t *testing.T) { } fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withMCP(mockMCP), @@ -1947,7 +1947,7 @@ func TestAnthropicToolChoiceParallelDisabled(t *testing.T) { require.Equal(t, http.StatusOK, resp.StatusCode) // Verify tool_choice in the upstream request. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) var receivedRequest map[string]any require.NoError(t, json.Unmarshal(received[0].Body, &receivedRequest)) @@ -2087,7 +2087,7 @@ func TestChatCompletionsParallelToolCallsDisabled(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) var opts []bridgeOption if tc.withInjectedTools { @@ -2112,7 +2112,7 @@ func TestChatCompletionsParallelToolCallsDisabled(t *testing.T) { _, err = io.ReadAll(resp.Body) require.NoError(t, err) - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) var upstreamReq map[string]any @@ -2142,7 +2142,7 @@ func TestThinkingAdaptiveIsPreserved(t *testing.T) { t.Cleanup(cancel) // Create a mock server that captures the request body sent upstream. - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -2160,7 +2160,7 @@ func TestThinkingAdaptiveIsPreserved(t *testing.T) { require.NoError(t, err) // Verify the thinking field was preserved in the upstream request. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) assert.Equal(t, "adaptive", gjson.GetBytes(received[0].Body, "thinking.type").Str) }) @@ -2207,7 +2207,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) // Set environment variables that the SDK would automatically read. // These should NOT leak into upstream requests. @@ -2223,7 +2223,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) { require.Equal(t, http.StatusOK, resp.StatusCode) // Verify that environment values did not leak. - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) require.Empty(t, received[0].Header.Get(tc.headerName)) }) @@ -2319,7 +2319,7 @@ func TestActorHeaders(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) metadataKey := "Username" bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, @@ -2341,7 +2341,7 @@ func TestActorHeaders(t *testing.T) { _, err = io.ReadAll(resp.Body) require.NoError(t, err) - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.NotEmpty(t, received) receivedHeaders := received[0].Header diff --git a/aibridge/internal/integrationtest/metrics_internal_test.go b/aibridge/internal/integrationtest/metrics_internal_test.go index dd3c010607..314c2d97c4 100644 --- a/aibridge/internal/integrationtest/metrics_internal_test.go +++ b/aibridge/internal/integrationtest/metrics_internal_test.go @@ -147,7 +147,7 @@ func TestMetrics_Interception(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) upstream.AllowOverflow = tc.allowOverflow m := aibridge.NewMetrics(prometheus.NewRegistry()) @@ -256,7 +256,7 @@ func TestMetrics_PromptCount(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.OaiChatSimple) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) m := aibridge.NewMetrics(prometheus.NewRegistry()) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, @@ -342,7 +342,7 @@ func TestMetrics_TokenUseCount(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) m := aibridge.NewMetrics(prometheus.NewRegistry()) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, @@ -383,7 +383,7 @@ func TestMetrics_NonInjectedToolUseCount(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, fixtures.OaiChatSingleBuiltinTool) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) m := aibridge.NewMetrics(prometheus.NewRegistry()) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, @@ -410,7 +410,7 @@ func TestMetrics_InjectedToolUseCount(t *testing.T) { // First request returns the tool invocation, the second returns the mocked response to the tool result. fix := fixtures.Parse(t, fixtures.AntSingleInjectedTool) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix), newFixtureToolResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix), testutil.NewFixtureToolResponse(fix)) m := aibridge.NewMetrics(prometheus.NewRegistry()) diff --git a/aibridge/internal/integrationtest/responses_internal_test.go b/aibridge/internal/integrationtest/responses_internal_test.go index 906f817500..4a6f2e30c4 100644 --- a/aibridge/internal/integrationtest/responses_internal_test.go +++ b/aibridge/internal/integrationtest/responses_internal_test.go @@ -328,7 +328,7 @@ func TestResponsesOutputMatchesUpstream(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) @@ -544,7 +544,7 @@ func TestResponsesParallelToolsOverwritten(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture[i]) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) var opts []bridgeOption if tc.withInjectedTools { @@ -567,7 +567,7 @@ func TestResponsesParallelToolsOverwritten(t *testing.T) { _, err = io.ReadAll(resp.Body) require.NoError(t, err) - received := upstream.receivedRequests() + received := upstream.ReceivedRequests() require.Len(t, received, 1) var upstreamReq map[string]any @@ -873,7 +873,7 @@ func TestResponsesInjectedTool(t *testing.T) { // Setup mock server for multi-turn interaction. // First request → tool call response, second → tool response. fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix), newFixtureToolResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix), testutil.NewFixtureToolResponse(fix)) // Setup MCP server proxies (with mock tools). mockMCP := setupMCPForTest(t, defaultTracer) @@ -1023,7 +1023,7 @@ func TestResponsesModelThoughts(t *testing.T) { t.Cleanup(cancel) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL) diff --git a/aibridge/internal/integrationtest/setupbridge.go b/aibridge/internal/integrationtest/setupbridge.go index 9674640b06..e63f554a00 100644 --- a/aibridge/internal/integrationtest/setupbridge.go +++ b/aibridge/internal/integrationtest/setupbridge.go @@ -216,10 +216,10 @@ func setupInjectedToolTest( // Setup mock server for multi-turn interaction. // First request → tool call response // Second request → final response. - firstResp := newFixtureResponse(fix) - toolResp := newFixtureToolResponse(fix) + firstResp := testutil.NewFixtureResponse(fix) + toolResp := testutil.NewFixtureToolResponse(fix) toolResp.OnRequest = toolRequestValidatorFn - upstream := newMockUpstream(ctx, t, firstResp, toolResp) + upstream := testutil.NewMockUpstream(ctx, t, firstResp, toolResp) mockMCP := setupMCPForTest(t, tracer) diff --git a/aibridge/internal/integrationtest/trace_internal_test.go b/aibridge/internal/integrationtest/trace_internal_test.go index c22635e9a9..baf7a5a3ae 100644 --- a/aibridge/internal/integrationtest/trace_internal_test.go +++ b/aibridge/internal/integrationtest/trace_internal_test.go @@ -147,7 +147,7 @@ func TestTraceAnthropic(t *testing.T) { sr, tracer := setupTracer(t) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) opts := []bridgeOption{ withTracer(tracer), @@ -263,7 +263,7 @@ func TestTraceAnthropicErr(t *testing.T) { sr, tracer := setupTracer(t) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) opts := []bridgeOption{ withTracer(tracer), @@ -552,7 +552,7 @@ func TestTraceOpenAI(t *testing.T) { sr, tracer := setupTracer(t) fix := fixtures.Parse(t, tc.fixture) - upstream := newMockUpstream(ctx, t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withTracer(tracer), ) @@ -711,7 +711,7 @@ func TestTraceOpenAIErr(t *testing.T) { fix := fixtures.Parse(t, tc.fixture) - mockAPI := newMockUpstream(ctx, t, newFixtureResponse(fix)) + mockAPI := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) mockAPI.AllowOverflow = tc.allowOverflow bridgeServer := newBridgeTestServer(ctx, t, mockAPI.URL, withTracer(tracer), @@ -753,7 +753,7 @@ func TestTracePassthrough(t *testing.T) { fix := fixtures.Parse(t, fixtures.OaiChatFallthrough) - upstream := newMockUpstream(t.Context(), t, newFixtureResponse(fix)) + upstream := testutil.NewMockUpstream(t.Context(), t, testutil.NewFixtureResponse(fix)) sr, tracer := setupTracer(t) diff --git a/aibridge/internal/testutil/mockserverproxier.go b/aibridge/internal/testutil/mockserverproxier.go new file mode 100644 index 0000000000..04f78330b3 --- /dev/null +++ b/aibridge/internal/testutil/mockserverproxier.go @@ -0,0 +1,47 @@ +package testutil + +import ( + "context" + + mcpgo "github.com/mark3labs/mcp-go/mcp" + + "github.com/coder/coder/v2/aibridge/mcp" +) + +// MockServerProxier is a test [mcp.ServerProxier] that injects a fixed set of +// tools. +type MockServerProxier struct { + Tools []*mcp.Tool +} + +func (*MockServerProxier) Init(context.Context) error { + return nil +} + +func (*MockServerProxier) Shutdown(context.Context) error { + return nil +} + +func (m *MockServerProxier) ListTools() []*mcp.Tool { + return m.Tools +} + +func (m *MockServerProxier) GetTool(id string) *mcp.Tool { + for _, t := range m.Tools { + if t.ID == id { + return t + } + } + return nil +} + +func (*MockServerProxier) CallTool(context.Context, string, any) (*mcpgo.CallToolResult, error) { + return nil, nil //nolint:nilnil // mock: no-op implementation +} + +// StubToolCaller is a minimal tool client that returns a fixed text result. +type StubToolCaller struct{} + +func (StubToolCaller) CallTool(_ context.Context, _ mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + return mcpgo.NewToolResultText("tool result"), nil +} diff --git a/aibridge/internal/integrationtest/mockupstream.go b/aibridge/internal/testutil/mockupstream.go similarity index 81% rename from aibridge/internal/integrationtest/mockupstream.go rename to aibridge/internal/testutil/mockupstream.go index cbef82047f..242bfde345 100644 --- a/aibridge/internal/integrationtest/mockupstream.go +++ b/aibridge/internal/testutil/mockupstream.go @@ -1,4 +1,4 @@ -package integrationtest +package testutil import ( "bufio" @@ -25,10 +25,10 @@ import ( "github.com/coder/coder/v2/aibridge/intercept/eventstream" ) -// upstreamResponse defines a single response that mockUpstream will replay -// for one incoming request. Use [newFixtureResponse] or [newFixtureToolResponse] to +// UpstreamResponse defines a single response that MockUpstream will replay +// for one incoming request. Use [NewFixtureResponse] or [NewFixtureToolResponse] to // construct one from a parsed txtar archive. -type upstreamResponse struct { +type UpstreamResponse struct { Streaming []byte // returned when the request has "stream": true. Blocking []byte // returned for non-streaming requests. @@ -37,11 +37,11 @@ type upstreamResponse struct { OnRequest func(r *http.Request, body []byte) } -// newFixtureResponse creates an upstreamResponse from a parsed fixture archive. +// NewFixtureResponse creates an UpstreamResponse from a parsed fixture archive. // It reads whichever of 'streaming' and 'non-streaming' sections exist; // not every fixture has both (e.g. error fixtures may only define one). -func newFixtureResponse(fix fixtures.Fixture) upstreamResponse { - var resp upstreamResponse +func NewFixtureResponse(fix fixtures.Fixture) UpstreamResponse { + var resp UpstreamResponse if fix.Has(fixtures.SectionStreaming) { resp.Streaming = fix.Streaming() } @@ -51,11 +51,11 @@ func newFixtureResponse(fix fixtures.Fixture) upstreamResponse { return resp } -// newFixtureToolResponse creates an upstreamResponse from the tool-call fixture files. +// NewFixtureToolResponse creates an UpstreamResponse from the tool-call fixture files. // It reads whichever of 'streaming/tool-call' and 'non-streaming/tool-call' // sections exist. -func newFixtureToolResponse(fix fixtures.Fixture) upstreamResponse { - var resp upstreamResponse +func NewFixtureToolResponse(fix fixtures.Fixture) UpstreamResponse { + var resp UpstreamResponse if fix.Has(fixtures.SectionStreamingToolCall) { resp.Streaming = fix.StreamingToolCall() } @@ -65,31 +65,36 @@ func newFixtureToolResponse(fix fixtures.Fixture) upstreamResponse { return resp } -// newErrorResponse returns an upstreamResponse that replays a raw HTTP error -// response with the given status code. Used to drive iteration-N error paths -// from inside a multi-call mockUpstream scripted-response list. -func newErrorResponse(status int) upstreamResponse { +// NewErrorResponse returns an UpstreamResponse that replays a raw HTTP error +// response with the given status code and optional Retry-After header. SDK +// auto-retries are disabled via x-should-retry. +func NewErrorResponse(status int, retryAfter string) UpstreamResponse { body := fmt.Sprintf(`{"error":{"message":%q}}`, http.StatusText(status)) + raw := fmt.Sprintf("HTTP/1.1 %d %s\r\n", status, http.StatusText(status)) + if retryAfter != "" { + raw += fmt.Sprintf("Retry-After: %s\r\n", retryAfter) + } raw += "x-should-retry: false\r\n" raw += "Content-Type: application/json\r\n" raw += fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body) + rawBytes := []byte(raw) - return upstreamResponse{Streaming: rawBytes, Blocking: rawBytes} + return UpstreamResponse{Streaming: rawBytes, Blocking: rawBytes} } -// receivedRequest captures the details of a single request handled by mockUpstream. -type receivedRequest struct { +// ReceivedRequest captures the details of a single request handled by MockUpstream. +type ReceivedRequest struct { Method string Path string Header http.Header Body []byte } -// mockUpstream replays txtar fixture responses, validates incoming request +// MockUpstream replays txtar fixture responses, validates incoming request // bodies, and counts calls. It stands in for a real AI provider API // (Anthropic, OpenAI) during integration tests. -type mockUpstream struct { +type MockUpstream struct { *httptest.Server // Calls is incremented atomically on every request. @@ -107,31 +112,31 @@ type mockUpstream struct { AllowOverflow bool mu sync.Mutex - requests []receivedRequest + requests []ReceivedRequest t *testing.T - responses []upstreamResponse + responses []UpstreamResponse } -// receivedRequests returns a copy of all requests received so far. -func (ms *mockUpstream) receivedRequests() []receivedRequest { +// ReceivedRequests returns a copy of all requests received so far. +func (ms *MockUpstream) ReceivedRequests() []ReceivedRequest { ms.mu.Lock() defer ms.mu.Unlock() - return append([]receivedRequest(nil), ms.requests...) + return append([]ReceivedRequest(nil), ms.requests...) } -// newMockUpstream creates a started httptest.Server that replays fixture +// NewMockUpstream creates a started httptest.Server that replays fixture // responses. Responses are returned in order: first call → first response. // The test fails if the number of requests doesn't match the number of // responses (when AllowOverflow is not set, default). // -// srv := newMockUpstream(ctx, t, newFixtureResponse(fix)) // simple -// srv := newMockUpstream(ctx, t, newFixtureResponse(fix), newFixtureToolResponse(fix)) // multi-turn -func newMockUpstream(ctx context.Context, t *testing.T, responses ...upstreamResponse) *mockUpstream { +// srv := NewMockUpstream(ctx, t, NewFixtureResponse(fix)) // simple +// srv := NewMockUpstream(ctx, t, NewFixtureResponse(fix), NewFixtureToolResponse(fix)) // multi-turn +func NewMockUpstream(ctx context.Context, t *testing.T, responses ...UpstreamResponse) *MockUpstream { t.Helper() - require.NotEmpty(t, responses, "at least one upstreamResponse required") + require.NotEmpty(t, responses, "at least one UpstreamResponse required") - ms := &mockUpstream{ + ms := &MockUpstream{ t: t, responses: responses, } @@ -156,7 +161,7 @@ func newMockUpstream(ctx context.Context, t *testing.T, responses ...upstreamRes return ms } -func (ms *mockUpstream) handle(w http.ResponseWriter, r *http.Request) { +func (ms *MockUpstream) handle(w http.ResponseWriter, r *http.Request) { call := int(ms.Calls.Add(1) - 1) body, err := io.ReadAll(r.Body) @@ -164,7 +169,7 @@ func (ms *mockUpstream) handle(w http.ResponseWriter, r *http.Request) { require.NoError(ms.t, err) ms.mu.Lock() - ms.requests = append(ms.requests, receivedRequest{ + ms.requests = append(ms.requests, ReceivedRequest{ Method: r.Method, Path: r.URL.Path, Header: r.Header.Clone(), @@ -201,7 +206,7 @@ func (ms *mockUpstream) handle(w http.ResponseWriter, r *http.Request) { _, _ = w.Write(resp.Blocking) } -func (ms *mockUpstream) responseForCall(call int) upstreamResponse { +func (ms *MockUpstream) responseForCall(call int) UpstreamResponse { if call >= len(ms.responses) { if ms.AllowOverflow { return ms.responses[len(ms.responses)-1] @@ -218,7 +223,7 @@ func isStreaming(body []byte, urlPath string) bool { return gjson.GetBytes(body, "stream").Bool() || strings.HasSuffix(urlPath, "invoke-with-response-stream") } -func (ms *mockUpstream) writeSSE(w http.ResponseWriter, data []byte) { +func (ms *MockUpstream) writeSSE(w http.ResponseWriter, data []byte) { ms.t.Helper() w.Header().Set("Content-Type", "text/event-stream") @@ -257,7 +262,7 @@ func isRawHTTPResponse(data []byte) bool { // writeRawHTTPResponse parses data as a complete HTTP response and replays it, // copying the status code, headers, and body to w. This supports error fixtures // that contain full HTTP responses (e.g. "HTTP/2.0 400 Bad Request\r\n..."). -func (ms *mockUpstream) writeRawHTTPResponse(w http.ResponseWriter, r *http.Request, data []byte) { +func (ms *MockUpstream) writeRawHTTPResponse(w http.ResponseWriter, r *http.Request, data []byte) { ms.t.Helper() resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(data)), r)