Review tweaks

Author: mtodor

cmd/stackrox-mcp/main.go

@@ -8,13 +8,13 @@ import (
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
-	flag "github.com/spf13/pflag"
+	"github.com/spf13/pflag"
 
 	"github.com/stackrox/stackrox-mcp/internal/config"
 	"github.com/stackrox/stackrox-mcp/internal/server"
 	"github.com/stackrox/stackrox-mcp/internal/toolsets"
-	configtools "github.com/stackrox/stackrox-mcp/internal/toolsets/config"
-	"github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability"
+	toolsetConfig "github.com/stackrox/stackrox-mcp/internal/toolsets/config"
+	toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability"
 )
 
 func setupLogging() {
@@ -24,27 +24,24 @@ func setupLogging() {
 // getToolsets initializes and returns all available toolsets
 func getToolsets(cfg *config.Config) []toolsets.Toolset {
 	return []toolsets.Toolset{
-		configtools.NewToolset(cfg),
-		vulnerability.NewToolset(cfg),
+		toolsetConfig.NewToolset(cfg),
+		toolsetVulnerability.NewToolset(cfg),
 	}
 }
 
 func main() {
 	setupLogging()
 
-	configPath := flag.String("config", "", "Path to configuration file (optional)")
-	flag.Parse()
+	configPath := pflag.String("config", "", "Path to configuration file (optional)")
+	pflag.Parse()
 
 	cfg, err := config.LoadConfig(*configPath)
 	if err != nil {
 		log.Fatal().Err(err).Msg("Failed to load configuration")
 	}
 	log.Info().Interface("config", cfg).Msg("Configuration loaded successfully")
 
-	// Create registry with all toolsets
 	registry := toolsets.NewRegistry(cfg, getToolsets(cfg))
-
-	// Create MCP server
 	srv := server.NewServer(cfg, registry)
 
 	// Set up context with signal handling for graceful shutdown
@@ -60,7 +57,6 @@ func main() {
 		cancel()
 	}()
 
-	// Start the server
 	log.Info().Msg("Starting StackRox MCP server")
 	if err := srv.Start(ctx); err != nil {
 		log.Fatal().Err(err).Msg("Server error")

cmd/stackrox-mcp/main_test.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"syscall"
 	"testing"
 	"time"
 
@@ -54,12 +53,11 @@ func TestGetToolsets(t *testing.T) {
 		},
 	}
 
-	toolsets := getToolsets(cfg)
-
-	require.NotNil(t, toolsets)
-	assert.Len(t, toolsets, 2, "Should have 2 toolsets")
-	assert.Equal(t, "config_manager", toolsets[0].GetName())
-	assert.Equal(t, "vulnerability", toolsets[1].GetName())
+	allToolsets := getToolsets(cfg)
+	require.NotNil(t, allToolsets)
+	assert.Len(t, allToolsets, 2, "Should have 2 toolsets")
+	assert.Equal(t, "config_manager", allToolsets[0].GetName())
+	assert.Equal(t, "vulnerability", allToolsets[1].GetName())
 }
 
 func TestGracefulShutdown(t *testing.T) {
@@ -74,11 +72,8 @@ func TestGracefulShutdown(t *testing.T) {
 	// Use a different port to avoid conflicts
 	cfg.Server.Port = 9999
 
-	// Create registry and server
 	registry := toolsets.NewRegistry(cfg, getToolsets(cfg))
 	srv := server.NewServer(cfg, registry)
-
-	// Set up context with cancellation
 	ctx, cancel := context.WithCancel(context.Background())
 
 	// Start server in goroutine
@@ -102,7 +97,7 @@ func TestGracefulShutdown(t *testing.T) {
 	// Simulate shutdown signal by canceling context
 	cancel()
 
-	// Wait for server to shut down with timeout
+	// Wait for server to shut down
 	select {
 	case err := <-errChan:
 		// Server should shut down cleanly (either nil or context.Canceled)
@@ -113,64 +108,3 @@ func TestGracefulShutdown(t *testing.T) {
 		t.Fatal("Server did not shut down within timeout period")
 	}
 }
-
-func TestGracefulShutdown_WithSignal(t *testing.T) {
-	// Set up minimal valid config
-	assert.NoError(t, os.Setenv("STACKROX_MCP__TOOLS__VULNERABILITY__ENABLED", "true"))
-	defer func() { assert.NoError(t, os.Unsetenv("STACKROX_MCP__TOOLS__VULNERABILITY__ENABLED")) }()
-
-	cfg, err := config.LoadConfig("")
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
-
-	// Use a different port to avoid conflicts
-	cfg.Server.Port = 10000
-
-	// Create registry and server
-	registry := toolsets.NewRegistry(cfg, getToolsets(cfg))
-	srv := server.NewServer(cfg, registry)
-
-	// Set up context with cancellation and signal handling (like in main)
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	sigChan := make(chan os.Signal, 1)
-	// Note: We use a buffered channel and manually send to avoid OS signal complications in tests
-
-	go func() {
-		<-sigChan
-		cancel()
-	}()
-
-	// Start server in goroutine
-	errChan := make(chan error, 1)
-	go func() {
-		errChan <- srv.Start(ctx)
-	}()
-
-	// Wait for server to be ready by polling
-	serverURL := fmt.Sprintf("http://%s:%d", cfg.Server.Address, cfg.Server.Port)
-	err = waitForServerReady(serverURL, 3*time.Second)
-	require.NoError(t, err, "Server should start within timeout")
-
-	// Establish actual HTTP connection to verify server is responding
-	resp, err := http.Get(serverURL)
-	if err == nil {
-		_ = resp.Body.Close()
-	}
-	assert.NoError(t, err, "Should be able to establish HTTP connection to server")
-
-	// Simulate signal by sending to channel
-	sigChan <- syscall.SIGTERM
-
-	// Wait for server to shut down with timeout
-	select {
-	case err := <-errChan:
-		// Server should shut down cleanly
-		if err != nil && err != context.Canceled {
-			t.Errorf("Server returned unexpected error: %v", err)
-		}
-	case <-time.After(5 * time.Second):
-		t.Fatal("Server did not shut down within timeout period after signal")
-	}
-}

internal/config/config.go

@@ -109,7 +109,7 @@ func (c *Config) Validate() error {
 		return fmt.Errorf("server.address is required")
 	}
 
-	if c.Server.Port <= 0 || c.Server.Port > 65535 {
+	if c.Server.Port < 1 || c.Server.Port > 65535 {
 		return fmt.Errorf("server.port must be between 1 and 65535")
 	}
 

internal/server/server.go

@@ -32,7 +32,7 @@ func NewServer(cfg *config.Config, registry *toolsets.Registry) *Server {
 			Name:    "stackrox-mcp",
 			Version: version,
 		},
-		nil, // No custom ServerOptions needed for now
+		nil,
 	)
 
 	return &Server{
@@ -79,12 +79,11 @@ func (s *Server) registerTools() {
 func (s *Server) Start(ctx context.Context) error {
 	s.registerTools()
 
-	// Create Streamable HTTP handler
 	handler := mcp.NewStreamableHTTPHandler(
 		func(*http.Request) *mcp.Server {
 			return s.mcp
 		},
-		nil, // No custom StreamableHTTPOptions needed
+		nil,
 	)
 
 	addr := fmt.Sprintf("%s:%d", s.cfg.Server.Address, s.cfg.Server.Port)
@@ -96,7 +95,6 @@ func (s *Server) Start(ctx context.Context) error {
 	log.Info().
 		Str("address", s.cfg.Server.Address).
 		Int("port", s.cfg.Server.Port).
-		Str("url", fmt.Sprintf("http://%s", addr)).
 		Msg("Starting MCP HTTP server")
 
 	// Start server in a goroutine

internal/toolsets/config/toolset_test.go

@@ -28,7 +28,7 @@ func TestToolset_IsEnabled_True(t *testing.T) {
 
 	tools := toolset.GetTools()
 	require.NotEmpty(t, tools, "Should return tools when enabled")
-	assert.Len(t, tools, 1, "Should have list_clusters tool")
+	require.Len(t, tools, 1, "Should have list_clusters tool")
 	assert.Equal(t, "list_clusters", tools[0].GetName())
 }
 

internal/toolsets/registry_test.go

@@ -136,7 +136,7 @@ func TestRegistry_GetAllTools_OneToolsetDisabled(t *testing.T) {
 	tools := registry.GetAllTools()
 
 	require.NotEmpty(t, tools)
-	assert.Len(t, tools, 1, "Should only have tools from enabled toolset")
+	require.Len(t, tools, 1, "Should only have tools from enabled toolset")
 	assert.Equal(t, "enabled_tool", tools[0].Name)
 }
 
@@ -180,8 +180,15 @@ func TestRegistry_GetAllTools_FiltersReadWriteTools(t *testing.T) {
 
 	// Should only have 2 read-only tools
 	require.Len(t, tools, 2, "Should filter out read-write tools")
-	assert.Equal(t, "read_only_1", tools[0].Name)
-	assert.Equal(t, "read_only_2", tools[1].Name)
+
+	// Verify tools are present
+	toolNames := make(map[string]bool)
+	for _, tool := range tools {
+		toolNames[tool.Name] = true
+	}
+
+	assert.True(t, toolNames["read_only_1"], "Should have read-only tools")
+	assert.True(t, toolNames["read_only_2"], "Should have read-only tools")
 }
 
 func TestRegistry_GetToolsets(t *testing.T) {
@@ -195,12 +202,11 @@ func TestRegistry_GetToolsets(t *testing.T) {
 	})
 
 	toolsetList := []toolsets.Toolset{mockToolset1, mockToolset2}
-
 	registry := toolsets.NewRegistry(cfg, toolsetList)
 
 	retrievedToolsets := registry.GetToolsets()
 
-	assert.Len(t, retrievedToolsets, 2)
+	require.Len(t, retrievedToolsets, 2)
 	assert.Contains(t, retrievedToolsets, mockToolset1)
 	assert.Contains(t, retrievedToolsets, mockToolset2)
 }

internal/toolsets/toolset.go

@@ -21,16 +21,12 @@ type Tool interface {
 
 // Toolset represents a collection of related tools
 type Toolset interface {
-	// GetName returns the toolset name (e.g., "config_manager", "vulnerability")
+	// GetName returns the toolset name
 	GetName() string
 
 	// IsEnabled checks if this toolset is enabled in configuration
 	IsEnabled() bool
 
 	// GetTools returns available tools based on configuration
-	// Returns empty list if:
-	// - Toolset is not enabled
-	// - read_only_tools=true and toolset has no read-only tools
-	// Filtering by read_only and enabled happens internally
 	GetTools() []Tool
 }

internal/toolsets/vulnerability/toolset_test.go

@@ -37,7 +37,7 @@ func TestToolset_IsEnabled_True(t *testing.T) {
 
 	tools := toolset.GetTools()
 	require.NotEmpty(t, tools, "Should return tools when enabled")
-	assert.Len(t, tools, 1, "Should have list_cluster_cves tool")
+	require.Len(t, tools, 1, "Should have list_cluster_cves tool")
 	assert.Equal(t, "list_cluster_cves", tools[0].GetName())
 }