Ensure toolsets are configurable via env var

cmd/github-mcp-server/main.go

@@ -45,7 +45,15 @@ var (
 				stdlog.Fatal("Failed to initialize logger:", err)
 			}
 
-			enabledToolsets := viper.GetStringSlice("toolsets")
+			// If you're wondering why we're not using viper.GetStringSlice("toolsets"),
+			// it's because viper doesn't handle comma-separated values correctly for env
+			// vars when using GetStringSlice.
+			// https://github.com/spf13/viper/issues/380
+			var enabledToolsets []string
+			err = viper.UnmarshalKey("toolsets", &enabledToolsets)
+			if err != nil {
+				stdlog.Fatal("Failed to unmarshal toolsets:", err)
+			}
 
 			logCommands := viper.GetBool("enable-command-logging")
 			cfg := runConfig{

e2e/e2e_test.go

@@ -5,8 +5,11 @@ package e2e_test
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"os"
 	"os/exec"
+	"slices"
+	"sync"
 	"testing"
 	"time"
 
@@ -16,85 +19,190 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestE2E(t *testing.T) {
-	e2eServerToken := os.Getenv("GITHUB_MCP_SERVER_E2E_TOKEN")
-	if e2eServerToken == "" {
-		t.Fatalf("GITHUB_MCP_SERVER_E2E_TOKEN environment variable is not set")
+var (
+	// Shared variables and sync.Once instances to ensure one-time execution
+	getTokenOnce sync.Once
+	token        string
+
+	buildOnce  sync.Once
+	buildError error
+)
+
+// getE2EToken ensures the environment variable is checked only once and returns the token
+func getE2EToken(t *testing.T) string {
+	getTokenOnce.Do(func() {
+		token = os.Getenv("GITHUB_MCP_SERVER_E2E_TOKEN")
+		if token == "" {
+			t.Fatalf("GITHUB_MCP_SERVER_E2E_TOKEN environment variable is not set")
+		}
+	})
+	return token
+}
+
+// ensureDockerImageBuilt makes sure the Docker image is built only once across all tests
+func ensureDockerImageBuilt(t *testing.T) {
+	buildOnce.Do(func() {
+		t.Log("Building Docker image for e2e tests...")
+		cmd := exec.Command("docker", "build", "-t", "github/e2e-github-mcp-server", ".")
+		cmd.Dir = ".." // Run this in the context of the root, where the Dockerfile is located.
+		output, err := cmd.CombinedOutput()
+		buildError = err
+		if err != nil {
+			t.Logf("Docker build output: %s", string(output))
+		}
+	})
+
+	// Check if the build was successful
+	require.NoError(t, buildError, "expected to build Docker image successfully")
+}
+
+// ClientOpts holds configuration options for the MCP client setup
+type ClientOpts struct {
+	// Environment variables to set before starting the client
+	EnvVars map[string]string
+}
+
+// ClientOption defines a function type for configuring ClientOpts
+type ClientOption func(*ClientOpts)
+
+// WithEnvVars returns an option that adds environment variables to the client options
+func WithEnvVars(envVars map[string]string) ClientOption {
+	return func(opts *ClientOpts) {
+		opts.EnvVars = envVars
+	}
+}
+
+// setupMCPClient sets up the test environment and returns an initialized MCP client
+// It handles token retrieval, Docker image building, and applying the provided options
+func setupMCPClient(t *testing.T, options ...ClientOption) *mcpClient.Client {
+	// Get token and ensure Docker image is built
+	token := getE2EToken(t)
+	ensureDockerImageBuilt(t)
+
+	// Create and configure options
+	opts := &ClientOpts{
+		EnvVars: make(map[string]string),
 	}
 
-	// Build the Docker image for the MCP server.
-	buildDockerImage(t)
+	// Apply all options to configure the opts struct
+	for _, option := range options {
+		option(opts)
+	}
 
-	t.Setenv("GITHUB_PERSONAL_ACCESS_TOKEN", e2eServerToken) // The MCP Client merges the existing environment.
+	// Prepare Docker arguments
 	args := []string{
 		"docker",
 		"run",
 		"-i",
 		"--rm",
 		"-e",
-		"GITHUB_PERSONAL_ACCESS_TOKEN",
-		"github/e2e-github-mcp-server",
+		"GITHUB_PERSONAL_ACCESS_TOKEN", // Personal access token is all required
+	}
+
+	// Add all environment variables to the Docker arguments
+	for key := range opts.EnvVars {
+		args = append(args, "-e", key)
+	}
+
+	// Add the image name
+	args = append(args, "github/e2e-github-mcp-server")
+
+	// Construct the env vars for the MCP Client to execute docker with
+	dockerEnvVars := make([]string, 0, len(opts.EnvVars)+1)
+	dockerEnvVars = append(dockerEnvVars, fmt.Sprintf("GITHUB_PERSONAL_ACCESS_TOKEN=%s", token))
+	for key, value := range opts.EnvVars {
+		dockerEnvVars = append(dockerEnvVars, fmt.Sprintf("%s=%s", key, value))
 	}
+
+	// Create the client
 	t.Log("Starting Stdio MCP client...")
-	client, err := mcpClient.NewStdioMCPClient(args[0], []string{}, args[1:]...)
+	client, err := mcpClient.NewStdioMCPClient(args[0], dockerEnvVars, args[1:]...)
 	require.NoError(t, err, "expected to create client successfully")
+	t.Cleanup(func() {
+		require.NoError(t, client.Close(), "expected to close client successfully")
+	})
 
-	t.Run("Initialize", func(t *testing.T) {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
+	// Initialize the client
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 
-		request := mcp.InitializeRequest{}
-		request.Params.ProtocolVersion = "2025-03-26"
-		request.Params.ClientInfo = mcp.Implementation{
-			Name:    "e2e-test-client",
-			Version: "0.0.1",
-		}
+	request := mcp.InitializeRequest{}
+	request.Params.ProtocolVersion = "2025-03-26"
+	request.Params.ClientInfo = mcp.Implementation{
+		Name:    "e2e-test-client",
+		Version: "0.0.1",
+	}
 
-		result, err := client.Initialize(ctx, request)
-		require.NoError(t, err, "expected to initialize successfully")
+	result, err := client.Initialize(ctx, request)
+	require.NoError(t, err, "failed to initialize client")
+	require.Equal(t, "github-mcp-server", result.ServerInfo.Name, "unexpected server name")
 
-		require.Equal(t, "github-mcp-server", result.ServerInfo.Name)
-	})
+	return client
+}
 
-	t.Run("CallTool get_me", func(t *testing.T) {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
+func TestGetMe(t *testing.T) {
+	t.Parallel()
 
-		// When we call the "get_me" tool
-		request := mcp.CallToolRequest{}
-		request.Params.Name = "get_me"
+	mcpClient := setupMCPClient(t)
 
-		response, err := client.CallTool(ctx, request)
-		require.NoError(t, err, "expected to call 'get_me' tool successfully")
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 
-		require.False(t, response.IsError, "expected result not to be an error")
-		require.Len(t, response.Content, 1, "expected content to have one item")
+	// When we call the "get_me" tool
+	request := mcp.CallToolRequest{}
+	request.Params.Name = "get_me"
 
-		textContent, ok := response.Content[0].(mcp.TextContent)
-		require.True(t, ok, "expected content to be of type TextContent")
+	response, err := mcpClient.CallTool(ctx, request)
+	require.NoError(t, err, "expected to call 'get_me' tool successfully")
 
-		var trimmedContent struct {
-			Login string `json:"login"`
-		}
-		err = json.Unmarshal([]byte(textContent.Text), &trimmedContent)
-		require.NoError(t, err, "expected to unmarshal text content successfully")
-
-		// Then the login in the response should match the login obtained via the same
-		// token using the GitHub API.
-		client := github.NewClient(nil).WithAuthToken(e2eServerToken)
-		user, _, err := client.Users.Get(context.Background(), "")
-		require.NoError(t, err, "expected to get user successfully")
-		require.Equal(t, trimmedContent.Login, *user.Login, "expected login to match")
-	})
+	require.False(t, response.IsError, "expected result not to be an error")
+	require.Len(t, response.Content, 1, "expected content to have one item")
+
+	textContent, ok := response.Content[0].(mcp.TextContent)
+	require.True(t, ok, "expected content to be of type TextContent")
+
+	var trimmedContent struct {
+		Login string `json:"login"`
+	}
+	err = json.Unmarshal([]byte(textContent.Text), &trimmedContent)
+	require.NoError(t, err, "expected to unmarshal text content successfully")
+
+	// Then the login in the response should match the login obtained via the same
+	// token using the GitHub API.
+	ghClient := github.NewClient(nil).WithAuthToken(getE2EToken(t))
+	user, _, err := ghClient.Users.Get(context.Background(), "")
+	require.NoError(t, err, "expected to get user successfully")
+	require.Equal(t, trimmedContent.Login, *user.Login, "expected login to match")
 
-	require.NoError(t, client.Close(), "expected to close client successfully")
 }
 
-func buildDockerImage(t *testing.T) {
-	t.Log("Building Docker image for e2e tests...")
+func TestToolsets(t *testing.T) {
+	t.Parallel()
+
+	mcpClient := setupMCPClient(
+		t,
+		WithEnvVars(map[string]string{
+			"GITHUB_TOOLSETS": "repos,issues",
+		}),
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	request := mcp.ListToolsRequest{}
+	response, err := mcpClient.ListTools(ctx, request)
+	require.NoError(t, err, "expected to list tools successfully")
+
+	// We could enumerate the tools here, but we'll need to expose that information
+	// declaratively in the MCP server, so for the moment let's just check the existence
+	// of an issue and repo tool, and the non-existence of a pull_request tool.
+	var toolsContains = func(expectedName string) bool {
+		return slices.ContainsFunc(response.Tools, func(tool mcp.Tool) bool {
+			return tool.Name == expectedName
+		})
+	}
 
-	cmd := exec.Command("docker", "build", "-t", "github/e2e-github-mcp-server", ".")
-	cmd.Dir = ".." // Run this in the context of the root, where the Dockerfile is located.
-	output, err := cmd.CombinedOutput()
-	require.NoError(t, err, "expected to build Docker image successfully, output: %s", string(output))
+	require.True(t, toolsContains("get_issue"), "expected to find 'get_issue' tool")
+	require.True(t, toolsContains("list_branches"), "expected to find 'list_branches' tool")
+	require.False(t, toolsContains("get_pull_request"), "expected not to find 'get_pull_request' tool")
 }