Skip to content

Commit cccb080

Browse files
johnstcnjohnstcn
committed
chore(mcp): fix test flakes
1 parent cc733ab commit cccb080

File tree

1 file changed

+34
-30
lines changed

1 file changed

+34
-30
lines changed

‎mcp/mcp_test.go

+34-30
Original file line numberDiff line numberDiff line change
@@ -77,15 +77,12 @@ func TestCoderTools(t *testing.T) {
7777
pty.WriteLine(ctr)
7878
_ = pty.ReadLine(ctx) // skip the echo
7979

80-
templates, err := memberClient.Templates(ctx, codersdk.TemplateFilter{})
80+
// Then: the response is a list of expected visible to the user.
81+
expected, err := memberClient.Templates(ctx, codersdk.TemplateFilter{})
8182
require.NoError(t, err)
82-
templatesJSON, err := json.Marshal(templates)
83-
require.NoError(t, err)
84-
85-
// Then: the response is a list of templates visible to the user.
86-
expected := makeJSONRPCTextResponse(t, string(templatesJSON))
87-
actual := pty.ReadLine(ctx)
88-
testutil.RequireJSONEq(t, expected, actual)
83+
actual := unmarshalFromCallToolResult[[]codersdk.Template](t, pty.ReadLine(ctx))
84+
require.Len(t, actual, 1)
85+
require.Equal(t, expected[0].ID, actual[0].ID)
8986
})
9087

9188
t.Run("coder_report_task", func(t *testing.T) {
@@ -111,20 +108,16 @@ func TestCoderTools(t *testing.T) {
111108

112109
t.Run("coder_whoami", func(t *testing.T) {
113110
// When: the coder_whoami tool is called
114-
me, err := memberClient.User(ctx, codersdk.Me)
115-
require.NoError(t, err)
116-
meJSON, err := json.Marshal(me)
117-
require.NoError(t, err)
118-
119111
ctr := makeJSONRPCRequest(t, "tools/call", "coder_whoami", map[string]any{})
120112

121113
pty.WriteLine(ctr)
122114
_ = pty.ReadLine(ctx) // skip the echo
123115

124116
// Then: the response is a valid JSON respresentation of the calling user.
125-
expected := makeJSONRPCTextResponse(t, string(meJSON))
126-
actual := pty.ReadLine(ctx)
127-
testutil.RequireJSONEq(t, expected, actual)
117+
expected, err := memberClient.User(ctx, codersdk.Me)
118+
require.NoError(t, err)
119+
actual := unmarshalFromCallToolResult[codersdk.User](t, pty.ReadLine(ctx))
120+
require.Equal(t, expected.ID, actual.ID)
128121
})
129122

130123
t.Run("coder_list_workspaces", func(t *testing.T) {
@@ -138,15 +131,10 @@ func TestCoderTools(t *testing.T) {
138131
pty.WriteLine(ctr)
139132
_ = pty.ReadLine(ctx) // skip the echo
140133

141-
ws, err := memberClient.Workspaces(ctx, codersdk.WorkspaceFilter{})
142-
require.NoError(t, err)
143-
wsJSON, err := json.Marshal(ws)
144-
require.NoError(t, err)
145-
146134
// Then: the response is a valid JSON respresentation of the calling user's workspaces.
147-
expected := makeJSONRPCTextResponse(t, string(wsJSON))
148-
actual := pty.ReadLine(ctx)
149-
testutil.RequireJSONEq(t, expected, actual)
135+
actual := unmarshalFromCallToolResult[codersdk.WorkspacesResponse](t, pty.ReadLine(ctx))
136+
require.Len(t, actual.Workspaces, 1, "expected 1 workspace")
137+
require.Equal(t, r.Workspace.ID, actual.Workspaces[0].ID, "expected the workspace to be the one we created in setup")
150138
})
151139

152140
t.Run("coder_get_workspace", func(t *testing.T) {
@@ -161,15 +149,12 @@ func TestCoderTools(t *testing.T) {
161149
pty.WriteLine(ctr)
162150
_ = pty.ReadLine(ctx) // skip the echo
163151

164-
ws, err := memberClient.Workspace(ctx, r.Workspace.ID)
165-
require.NoError(t, err)
166-
wsJSON, err := json.Marshal(ws)
152+
expected, err := memberClient.Workspace(ctx, r.Workspace.ID)
167153
require.NoError(t, err)
168154

169155
// Then: the response is a valid JSON respresentation of the workspace.
170-
expected := makeJSONRPCTextResponse(t, string(wsJSON))
171-
actual := pty.ReadLine(ctx)
172-
testutil.RequireJSONEq(t, expected, actual)
156+
actual := unmarshalFromCallToolResult[codersdk.Workspace](t, pty.ReadLine(ctx))
157+
require.Equal(t, expected.ID, actual.ID)
173158
})
174159

175160
// NOTE: this test runs after the list_workspaces tool is called.
@@ -322,6 +307,25 @@ func makeJSONRPCTextResponse(t *testing.T, text string) string {
322307
return string(bs)
323308
}
324309

310+
func unmarshalFromCallToolResult[T any](t *testing.T, raw string) T {
311+
t.Helper()
312+
313+
var resp map[string]any
314+
require.NoError(t, json.Unmarshal([]byte(raw), &resp), "failed to unmarshal JSON RPC response")
315+
res, ok := resp["result"].(map[string]any)
316+
require.True(t, ok, "expected a result field in the response")
317+
ct, ok := res["content"].([]any)
318+
require.True(t, ok, "expected a content field in the result")
319+
require.Len(t, ct, 1, "expected a single content item in the result")
320+
ct0, ok := ct[0].(map[string]any)
321+
require.True(t, ok, "expected a content item in the result")
322+
txt, ok := ct0["text"].(string)
323+
require.True(t, ok, "expected a text field in the content item")
324+
var actual T
325+
require.NoError(t, json.Unmarshal([]byte(txt), &actual), "failed to unmarshal content")
326+
return actual
327+
}
328+
325329
// startTestMCPServer is a helper function that starts a MCP server listening on
326330
// a pty. It is the responsibility of the caller to close the server.
327331
func startTestMCPServer(ctx context.Context, t testing.TB, stdin io.Reader, stdout io.Writer) (*server.MCPServer, func() error) {

0 commit comments

Comments
 (0)