test: split test client creation and tool retrieval into separate functions

Author: jochil

filesystemserver/inprocess_test.go

@@ -1,12 +1,9 @@
 package filesystemserver_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/mark3labs/mcp-filesystem-server/filesystemserver"
-	"github.com/mark3labs/mcp-go/client"
-	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -15,37 +12,9 @@ func TestInProcess(t *testing.T) {
 	fss, err := filesystemserver.NewFilesystemServer([]string{"."})
 	require.NoError(t, err)
 
-	mcpClient, err := client.NewInProcessClient(fss)
-	require.NoError(t, err)
-
-	err = mcpClient.Start(context.Background())
-	require.NoError(t, err)
-
-	// Initialize the client
-	initRequest := mcp.InitializeRequest{}
-	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
-	initRequest.Params.ClientInfo = mcp.Implementation{
-		Name:    "test-client",
-		Version: "1.0.0",
-	}
-	result, err := mcpClient.Initialize(context.Background(), initRequest)
-	require.NoError(t, err)
-	assert.Equal(t, "secure-filesystem-server", result.ServerInfo.Name)
-	assert.Equal(t, filesystemserver.Version, result.ServerInfo.Version)
+	mcpClient := startTestClient(t, fss)
 
 	// just check for a specific tool
-	toolListResult, err := mcpClient.ListTools(context.Background(), mcp.ListToolsRequest{})
-	require.NoError(t, err)
-	assert.NotEmpty(t, toolListResult.Tools)
-	found := false
-	for _, tool := range toolListResult.Tools {
-		if tool.Name == "read_file" {
-			found = true
-			break
-		}
-	}
-	assert.True(t, found, "read_file tool not found in the list of tools")
-
-	err = mcpClient.Close()
-	require.NoError(t, err)
+	tool := getTool(t, mcpClient, "read_file")
+	assert.NotNil(t, tool, "read_file tool not found in the list of tools")
 }

filesystemserver/utils_test.go

@@ -0,0 +1,51 @@
+package filesystemserver_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/mark3labs/mcp-filesystem-server/filesystemserver"
+	"github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func startTestClient(t *testing.T, fss *server.MCPServer) client.MCPClient {
+	t.Helper()
+
+	mcpClient, err := client.NewInProcessClient(fss)
+	require.NoError(t, err)
+	t.Cleanup(func() { mcpClient.Close() })
+
+	err = mcpClient.Start(context.Background())
+	require.NoError(t, err)
+
+	// Initialize the client
+	initRequest := mcp.InitializeRequest{}
+	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
+	initRequest.Params.ClientInfo = mcp.Implementation{
+		Name:    "test-client",
+		Version: "1.0.0",
+	}
+	result, err := mcpClient.Initialize(context.Background(), initRequest)
+	require.NoError(t, err)
+	assert.Equal(t, "secure-filesystem-server", result.ServerInfo.Name)
+	assert.Equal(t, filesystemserver.Version, result.ServerInfo.Version)
+
+	return mcpClient
+}
+
+func getTool(t *testing.T, mcpClient client.MCPClient, toolName string) *mcp.Tool {
+	result, err := mcpClient.ListTools(context.Background(), mcp.ListToolsRequest{})
+	require.NoError(t, err)
+	for _, tool := range result.Tools {
+		if tool.Name == toolName {
+			return &tool
+		}
+	}
+	require.Fail(t, "Tool not found", toolName)
+	return nil
+}