chore(mcp): fix test flakes (#17183)

Closes https://github.com/coder/internal/issues/547
This commit is contained in:
Cian Johnston
2025-04-01 11:28:47 +01:00
committed by GitHub
parent cc733aba71
commit e4cf18989c

View File

@ -77,15 +77,12 @@ func TestCoderTools(t *testing.T) {
pty.WriteLine(ctr) pty.WriteLine(ctr)
_ = pty.ReadLine(ctx) // skip the echo _ = pty.ReadLine(ctx) // skip the echo
templates, err := memberClient.Templates(ctx, codersdk.TemplateFilter{}) // Then: the response is a list of expected visible to the user.
expected, err := memberClient.Templates(ctx, codersdk.TemplateFilter{})
require.NoError(t, err) require.NoError(t, err)
templatesJSON, err := json.Marshal(templates) actual := unmarshalFromCallToolResult[[]codersdk.Template](t, pty.ReadLine(ctx))
require.NoError(t, err) require.Len(t, actual, 1)
require.Equal(t, expected[0].ID, actual[0].ID)
// Then: the response is a list of templates visible to the user.
expected := makeJSONRPCTextResponse(t, string(templatesJSON))
actual := pty.ReadLine(ctx)
testutil.RequireJSONEq(t, expected, actual)
}) })
t.Run("coder_report_task", func(t *testing.T) { t.Run("coder_report_task", func(t *testing.T) {
@ -111,20 +108,16 @@ func TestCoderTools(t *testing.T) {
t.Run("coder_whoami", func(t *testing.T) { t.Run("coder_whoami", func(t *testing.T) {
// When: the coder_whoami tool is called // When: the coder_whoami tool is called
me, err := memberClient.User(ctx, codersdk.Me)
require.NoError(t, err)
meJSON, err := json.Marshal(me)
require.NoError(t, err)
ctr := makeJSONRPCRequest(t, "tools/call", "coder_whoami", map[string]any{}) ctr := makeJSONRPCRequest(t, "tools/call", "coder_whoami", map[string]any{})
pty.WriteLine(ctr) pty.WriteLine(ctr)
_ = pty.ReadLine(ctx) // skip the echo _ = pty.ReadLine(ctx) // skip the echo
// Then: the response is a valid JSON respresentation of the calling user. // Then: the response is a valid JSON respresentation of the calling user.
expected := makeJSONRPCTextResponse(t, string(meJSON)) expected, err := memberClient.User(ctx, codersdk.Me)
actual := pty.ReadLine(ctx) require.NoError(t, err)
testutil.RequireJSONEq(t, expected, actual) actual := unmarshalFromCallToolResult[codersdk.User](t, pty.ReadLine(ctx))
require.Equal(t, expected.ID, actual.ID)
}) })
t.Run("coder_list_workspaces", func(t *testing.T) { t.Run("coder_list_workspaces", func(t *testing.T) {
@ -138,15 +131,10 @@ func TestCoderTools(t *testing.T) {
pty.WriteLine(ctr) pty.WriteLine(ctr)
_ = pty.ReadLine(ctx) // skip the echo _ = pty.ReadLine(ctx) // skip the echo
ws, err := memberClient.Workspaces(ctx, codersdk.WorkspaceFilter{})
require.NoError(t, err)
wsJSON, err := json.Marshal(ws)
require.NoError(t, err)
// Then: the response is a valid JSON respresentation of the calling user's workspaces. // Then: the response is a valid JSON respresentation of the calling user's workspaces.
expected := makeJSONRPCTextResponse(t, string(wsJSON)) actual := unmarshalFromCallToolResult[codersdk.WorkspacesResponse](t, pty.ReadLine(ctx))
actual := pty.ReadLine(ctx) require.Len(t, actual.Workspaces, 1, "expected 1 workspace")
testutil.RequireJSONEq(t, expected, actual) require.Equal(t, r.Workspace.ID, actual.Workspaces[0].ID, "expected the workspace to be the one we created in setup")
}) })
t.Run("coder_get_workspace", func(t *testing.T) { t.Run("coder_get_workspace", func(t *testing.T) {
@ -161,15 +149,12 @@ func TestCoderTools(t *testing.T) {
pty.WriteLine(ctr) pty.WriteLine(ctr)
_ = pty.ReadLine(ctx) // skip the echo _ = pty.ReadLine(ctx) // skip the echo
ws, err := memberClient.Workspace(ctx, r.Workspace.ID) expected, err := memberClient.Workspace(ctx, r.Workspace.ID)
require.NoError(t, err)
wsJSON, err := json.Marshal(ws)
require.NoError(t, err) require.NoError(t, err)
// Then: the response is a valid JSON respresentation of the workspace. // Then: the response is a valid JSON respresentation of the workspace.
expected := makeJSONRPCTextResponse(t, string(wsJSON)) actual := unmarshalFromCallToolResult[codersdk.Workspace](t, pty.ReadLine(ctx))
actual := pty.ReadLine(ctx) require.Equal(t, expected.ID, actual.ID)
testutil.RequireJSONEq(t, expected, actual)
}) })
// NOTE: this test runs after the list_workspaces tool is called. // NOTE: this test runs after the list_workspaces tool is called.
@ -322,6 +307,25 @@ func makeJSONRPCTextResponse(t *testing.T, text string) string {
return string(bs) return string(bs)
} }
func unmarshalFromCallToolResult[T any](t *testing.T, raw string) T {
t.Helper()
var resp map[string]any
require.NoError(t, json.Unmarshal([]byte(raw), &resp), "failed to unmarshal JSON RPC response")
res, ok := resp["result"].(map[string]any)
require.True(t, ok, "expected a result field in the response")
ct, ok := res["content"].([]any)
require.True(t, ok, "expected a content field in the result")
require.Len(t, ct, 1, "expected a single content item in the result")
ct0, ok := ct[0].(map[string]any)
require.True(t, ok, "expected a content item in the result")
txt, ok := ct0["text"].(string)
require.True(t, ok, "expected a text field in the content item")
var actual T
require.NoError(t, json.Unmarshal([]byte(txt), &actual), "failed to unmarshal content")
return actual
}
// startTestMCPServer is a helper function that starts a MCP server listening on // startTestMCPServer is a helper function that starts a MCP server listening on
// a pty. It is the responsibility of the caller to close the server. // a pty. It is the responsibility of the caller to close the server.
func startTestMCPServer(ctx context.Context, t testing.TB, stdin io.Reader, stdout io.Writer) (*server.MCPServer, func() error) { func startTestMCPServer(ctx context.Context, t testing.TB, stdin io.Reader, stdout io.Writer) (*server.MCPServer, func() error) {