From 0dc21a82d146f6d1bc8b745aff4485dd8191973e Mon Sep 17 00:00:00 2001
From: Marc Nuri <marc@marcnuri.com>
Date: Mon, 3 Nov 2025 15:51:24 +0100
Subject: [PATCH] test(mcp): refactor WatchKubeConfig tests to use testify and
 improve structure

Signed-off-by: Marc Nuri <marc@marcnuri.com>
---
 internal/test/mcp.go |  1 +
 pkg/mcp/mcp_test.go  | 68 +++++++++++++++++++++++++-------------------
 2 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/internal/test/mcp.go b/internal/test/mcp.go
index 5fa0d0a4..0b411b1d 100644
--- a/internal/test/mcp.go
+++ b/internal/test/mcp.go
@@ -23,6 +23,7 @@ func NewMcpClient(t *testing.T, mcpHttpServer http.Handler, options ...transport
 	var err error
 	ret := &McpClient{ctx: t.Context()}
 	ret.testServer = httptest.NewServer(mcpHttpServer)
+	options = append(options, transport.WithContinuousListening())
 	ret.Client, err = client.NewStreamableHttpClient(ret.testServer.URL+"/mcp", options...)
 	require.NoError(t, err, "Expected no error creating MCP client")
 	err = ret.Start(t.Context())
diff --git a/pkg/mcp/mcp_test.go b/pkg/mcp/mcp_test.go
index 484d8b59..9dca88e4 100644
--- a/pkg/mcp/mcp_test.go
+++ b/pkg/mcp/mcp_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"net/http"
 	"os"
-	"path/filepath"
 	"runtime"
 	"testing"
 	"time"
@@ -15,38 +14,47 @@ import (
 	"github.com/stretchr/testify/suite"
 )
 
-func TestWatchKubeConfig(t *testing.T) {
+type WatchKubeConfigSuite struct {
+	BaseMcpSuite
+}
+
+func (s *WatchKubeConfigSuite) SetupTest() {
+	s.BaseMcpSuite.SetupTest()
+	kubeconfig := test.KubeConfigFake()
+	s.Cfg.KubeConfig = test.KubeconfigFile(s.T(), kubeconfig)
+}
+
+func (s *WatchKubeConfigSuite) TestNotifiesToolsChange() {
 	if runtime.GOOS != "linux" && runtime.GOOS != "darwin" {
-		t.Skip("Skipping test on non-Unix-like platforms")
+		s.T().Skip("Skipping test on non-Unix-like platforms")
 	}
-	testCase(t, func(c *mcpContext) {
-		// Given
-		withTimeout, cancel := context.WithTimeout(c.ctx, 5*time.Second)
-		defer cancel()
-		var notification *mcp.JSONRPCNotification
-		c.mcpClient.OnNotification(func(n mcp.JSONRPCNotification) {
-			notification = &n
-		})
-		// When
-		f, _ := os.OpenFile(filepath.Join(c.tempDir, "config"), os.O_APPEND|os.O_WRONLY, 0644)
-		_, _ = f.WriteString("\n")
-		for notification == nil {
-			select {
-			case <-withTimeout.Done():
-			default:
-				time.Sleep(100 * time.Millisecond)
-			}
-		}
-		// Then
-		t.Run("WatchKubeConfig notifies tools change", func(t *testing.T) {
-			if notification == nil {
-				t.Fatalf("WatchKubeConfig did not notify")
-			}
-			if notification.Method != "notifications/tools/list_changed" {
-				t.Fatalf("WatchKubeConfig did not notify tools change, got %s", notification.Method)
-			}
-		})
+	// Given
+	s.InitMcpClient()
+	withTimeout, cancel := context.WithTimeout(s.T().Context(), 5*time.Second)
+	defer cancel()
+	var notification *mcp.JSONRPCNotification
+	s.OnNotification(func(n mcp.JSONRPCNotification) {
+		notification = &n
 	})
+	// When
+	f, _ := os.OpenFile(s.Cfg.KubeConfig, os.O_APPEND|os.O_WRONLY, 0644)
+	_, _ = f.WriteString("\n")
+	_ = f.Close()
+	for notification == nil {
+		select {
+		case <-withTimeout.Done():
+			s.FailNow("timeout waiting for WatchKubeConfig notification")
+		default:
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+	// Then
+	s.NotNil(notification, "WatchKubeConfig did not notify")
+	s.Equal("notifications/tools/list_changed", notification.Method, "WatchKubeConfig did not notify tools change")
+}
+
+func TestWatchKubeConfig(t *testing.T) {
+	suite.Run(t, new(WatchKubeConfigSuite))
 }
 
 type McpHeadersSuite struct {