mcp: fix more issues found in unit tests

Author: fgrosse

mcp/cmd_test.go

@@ -36,8 +36,6 @@ func SayHi(ctx context.Context, req *mcp.CallToolRequest, args SayHiParams) (*mc
 }
 
 func TestMain(m *testing.M) {
-	goleak.VerifyTestMain(m)
-
 	// If the runAsServer variable is set, execute the relevant serverFunc
 	// instead of running tests (aka the fork and exec trick).
 	if name := os.Getenv(runAsServer); name != "" {
@@ -49,6 +47,8 @@ func TestMain(m *testing.M) {
 		run()
 		return
 	}
+
+	goleak.VerifyTestMain(m)
 	os.Exit(m.Run())
 }
 
@@ -127,36 +127,35 @@ func TestServerInterrupt(t *testing.T) {
 	}
 	requireExec(t)
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
+	t.Log("Starting server command")
 	cmd := createServerCommand(t, "default")
 
 	client := mcp.NewClient(testImpl, nil)
+	t.Log("Connecting to server")
+
+	ctx := context.Background()
 	session, err := client.Connect(ctx, &mcp.CommandTransport{Command: cmd}, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	t.Cleanup(func() { session.Close() })
 
-	// get a signal when the server process exits
-	onExit := make(chan struct{})
-	go func() {
-		cmd.Process.Wait()
-		close(onExit)
-	}()
+	_, err = session.ListTools(ctx, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	// send a signal to the server process to terminate it
+	t.Log("Send a signal to the server process to terminate it")
 	if err := cmd.Process.Signal(os.Interrupt); err != nil {
 		t.Fatal(err)
 	}
 
-	// wait for the server to exit
-	// TODO: use synctest when available
-	select {
-	case <-time.After(5 * time.Second):
-		t.Fatal("server did not exit after SIGINT")
-	case <-onExit:
+	t.Log("Closing client session so server can exit immediately")
+	session.Close()
+
+	t.Log("Wait for process to terminate after interrupt signal")
+	_, err = cmd.Process.Wait()
+	if err == nil {
+		t.Errorf("unexpected error: %v", err)
 	}
 }
 

mcp/streamable_test.go

@@ -316,7 +316,7 @@ func testClientReplay(t *testing.T, test clientReplayTest) {
 		})
 
 	realServer := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil))
-	defer realServer.Close()
+	t.Cleanup(func() { realServer.Close() })
 	realServerURL, err := url.Parse(realServer.URL)
 	if err != nil {
 		t.Fatalf("Failed to parse real server URL: %v", err)