Skip to content

refactor(aibridge): move shared mock helpers to testutil #25999

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ssncferreira wants to merge 1 commit into
base: main
Choose a base branch
from
+163 −236
Draft

refactor(aibridge): move shared mock helpers to testutil #25999

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 3 additions & 43 deletions aibridge/intercept/chatcompletions/blocking_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package chatcompletions

import (
"context"
"io"
"net/http"
"net/http/httptest"
Expand All @@ -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"
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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
}
6 changes: 3 additions & 3 deletions aibridge/intercept/chatcompletions/streaming_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
68 changes: 19 additions & 49 deletions aibridge/intercept/messages/base_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,14 @@ 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"
"golang.org/x/xerrors"

"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"
Expand Down Expand Up @@ -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(),
}

Expand All @@ -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"},
},
},
Expand Down Expand Up @@ -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"},
},
},
Expand Down Expand Up @@ -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"},
},
},
Expand Down Expand Up @@ -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(),
}

Expand All @@ -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(),
}
Expand All @@ -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(),
}
Expand All @@ -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(),
}
Expand All @@ -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(),
}
Expand All @@ -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(),
}
Expand Down Expand Up @@ -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()

Expand Down
6 changes: 3 additions & 3 deletions aibridge/intercept/messages/blocking_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
16 changes: 3 additions & 13 deletions aibridge/intercept/messages/streaming_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package messages

import (
"context"
"io"
"net/http"
"net/http/httptest"
Expand All @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -541,10 +531,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",
Expand Down
46 changes: 3 additions & 43 deletions aibridge/intercept/responses/blocking_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package responses

import (
"context"
"io"
"net/http"
"net/http/httptest"
Expand All @@ -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"
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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
}
6 changes: 3 additions & 3 deletions aibridge/intercept/responses/streaming_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion aibridge/internal/integrationtest/apidump_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading
Loading