From 1d35a3dffff5cdbf37459687e84b2655659a651d Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Mon, 2 Mar 2026 11:58:54 +0000 Subject: [PATCH 01/19] ROX-31495: wiremock for central (#32) Signed-off-by: Tomasz Janiszewski Co-authored-by: Claude Sonnet 4.5 Co-authored-by: Mladen Todorovic Add Go integration tests for MCP server with WireMock Implements integration tests that verify MCP server functionality using stdio transport and WireMock as a mock StackRox Central backend. **Key Changes:** - Created integration test suite in `integration/` with build tag - Implemented stdio-based MCP client in `internal/testutil/mcp.go` - Added WireMock readiness check in `internal/testutil/wiremock.go` - Added test fixtures for expected WireMock response data - Updated Makefile with integration test targets - Updated GitHub Actions workflow to run integration tests in CI **Test Coverage:** - MCP protocol (initialize, list tools) - Tool invocations (list_clusters, get_deployments_for_cve, etc.) - Error handling (missing parameters) - Success and error scenarios Tests use stdio transport for simplicity and better control over the MCP server lifecycle. Each test starts a fresh MCP server subprocess and communicates via JSON-RPC over stdin/stdout. TODO: Fix WireMock request matching for CVE-2021-44228 deployment query Co-Authored-By: Claude Sonnet 4.5 fix: Fix WireMock JSONPath patterns for deployment CVE queries Apply the same JSONPath fix from commit 01f58ab to deployments.json. The original mappings used $.query[?(@.query =~ ...)] which looked for a nested array structure, but gRPC protobuf-to-JSON conversion creates a simple object with a 'query' field (lowercase). Changed all deployment CVE mappings from: $.query[?(@.query =~ /.*CVE-XXX.*/)] to: $[?(@.query =~ /.*CVE-XXX.*/)] This fixes the TestIntegration_GetDeploymentsForCVE_Log4Shell test which was failing because WireMock couldn't match the gRPC requests. Co-Authored-By: Claude Sonnet 4.5 test: Unskip Log4Shell integration test The WireMock JSONPath fix for deployments.json resolves the issue that was causing this test to fail. Co-Authored-By: Claude Sonnet 4.5 refactor: Simplify integration tests and remove manual protoc installation This commit implements two major simplifications to the test infrastructure: 1. **Eliminate TestMain Pre-compilation Pattern**: - Extract main() body into new internal/app package with Run() function - Tests now call app.Run() in-process via io.Pipe() instead of subprocess - Removes 60+ lines of build/setup code from integration_test.go - Enables full code coverage (previously main.go had 0% coverage) - Faster test execution (no binary compilation overhead) - Better debugging (direct function calls vs exec) 2. **Remove Manual Protoc Installation from GitHub Actions**: - Delete manual protoc 3.20.1 download from workflow - Rely on Makefile's automatic protoc 32.1 installation - Single source of truth for protoc version - Eliminates version mismatch between CI and local dev Changes: - Create internal/app/app.go with Run() function extracted from main() - Update server.Start() to accept optional stdin/stdout parameters - Refactor testutil.NewMCPClient() to use ServerRunFunc callback - Remove TestMain, buildMCPBinary, global vars from integration tests - Update all integration test functions to use createMCPClient() helper - Remove "Install protoc" step from .github/workflows/test.yml Benefits: - 43 net lines removed (-161 +118) - Better code coverage - Simpler test maintenance - Aligned with Go testing best practices Co-Authored-By: Claude Sonnet 4.5 fix: Add missing retry/timeout config to integration test setup The test config was missing RequestTimeout, MaxRetries, InitialBackoff, and MaxBackoff fields, causing the gRPC client to use zero values instead of the defaults (30s timeout, 3 retries). With MaxRetries=0, the retry loop never executed: `for attempt := range 0` This caused "Request failed after all retries ... attempts=0" warnings and empty responses from WireMock. Solution: Explicitly set the retry/timeout fields to match the defaults defined in internal/config/config.go. Co-Authored-By: Claude Sonnet 4.5 Signed-off-by: Tomasz Janiszewski Add Go integration tests for MCP server with WireMock Implements integration tests that verify MCP server functionality using stdio transport and WireMock as a mock StackRox Central backend. **Key Changes:** - Created integration test suite in `integration/` with build tag - Implemented stdio-based MCP client in `internal/testutil/mcp.go` - Added WireMock readiness check in `internal/testutil/wiremock.go` - Added test fixtures for expected WireMock response data - Updated Makefile with integration test targets - Updated GitHub Actions workflow to run integration tests in CI **Test Coverage:** - MCP protocol (initialize, list tools) - Tool invocations (list_clusters, get_deployments_for_cve, etc.) - Error handling (missing parameters) - Success and error scenarios Tests use stdio transport for simplicity and better control over the MCP server lifecycle. Each test starts a fresh MCP server subprocess and communicates via JSON-RPC over stdin/stdout. TODO: Fix WireMock request matching for CVE-2021-44228 deployment query Co-Authored-By: Claude Sonnet 4.5 fix: Fix WireMock JSONPath patterns for deployment CVE queries Apply the same JSONPath fix from commit 01f58ab to deployments.json. The original mappings used $.query[?(@.query =~ ...)] which looked for a nested array structure, but gRPC protobuf-to-JSON conversion creates a simple object with a 'query' field (lowercase). Changed all deployment CVE mappings from: $.query[?(@.query =~ /.*CVE-XXX.*/)] to: $[?(@.query =~ /.*CVE-XXX.*/)] This fixes the TestIntegration_GetDeploymentsForCVE_Log4Shell test which was failing because WireMock couldn't match the gRPC requests. Co-Authored-By: Claude Sonnet 4.5 test: Unskip Log4Shell integration test The WireMock JSONPath fix for deployments.json resolves the issue that was causing this test to fail. Co-Authored-By: Claude Sonnet 4.5 refactor: Simplify integration tests and remove manual protoc installation This commit implements two major simplifications to the test infrastructure: 1. **Eliminate TestMain Pre-compilation Pattern**: - Extract main() body into new internal/app package with Run() function - Tests now call app.Run() in-process via io.Pipe() instead of subprocess - Removes 60+ lines of build/setup code from integration_test.go - Enables full code coverage (previously main.go had 0% coverage) - Faster test execution (no binary compilation overhead) - Better debugging (direct function calls vs exec) 2. **Remove Manual Protoc Installation from GitHub Actions**: - Delete manual protoc 3.20.1 download from workflow - Rely on Makefile's automatic protoc 32.1 installation - Single source of truth for protoc version - Eliminates version mismatch between CI and local dev Changes: - Create internal/app/app.go with Run() function extracted from main() - Update server.Start() to accept optional stdin/stdout parameters - Refactor testutil.NewMCPClient() to use ServerRunFunc callback - Remove TestMain, buildMCPBinary, global vars from integration tests - Update all integration test functions to use createMCPClient() helper - Remove "Install protoc" step from .github/workflows/test.yml Benefits: - 43 net lines removed (-161 +118) - Better code coverage - Simpler test maintenance - Aligned with Go testing best practices Co-Authored-By: Claude Sonnet 4.5 fix: Add missing retry/timeout config to integration test setup The test config was missing RequestTimeout, MaxRetries, InitialBackoff, and MaxBackoff fields, causing the gRPC client to use zero values instead of the defaults (30s timeout, 3 retries). With MaxRetries=0, the retry loop never executed: `for attempt := range 0` This caused "Request failed after all retries ... attempts=0" warnings and empty responses from WireMock. Solution: Explicitly set the retry/timeout fields to match the defaults defined in internal/config/config.go. Co-Authored-By: Claude Sonnet 4.5 Signed-off-by: Tomasz Janiszewski cleanup Signed-off-by: Tomasz Janiszewski Refactor integration tests and migrate to official MCP SDK This commit refactors the integration test suite and replaces the custom MCP client implementation with the official MCP Go SDK. Changes: 1. Table-Driven Tests Refactoring: - Consolidated 5 individual test functions into 2 table-driven tests - TestIntegration_ToolCalls: 4 successful tool call scenarios - TestIntegration_ToolCallErrors: error handling scenarios - Reduced code duplication by ~50 lines - Added helper functions: setupInitializedClient, callToolAndGetResult 2. Removed TestMain: - Eliminated WireMock readiness check from TestMain - Removed unused imports (fmt, os) - Simplified test setup (13 lines removed) 3. Migrated to Official MCP Go SDK: - Replaced custom internal/testutil/mcp.go (202 lines) - Created internal/testutil/mcp_client.go (141 lines) using SDK - Uses official mcp.Client and mcp.ClientSession - Proper type-safe content handling with *mcp.TextContent - Better error handling (protocol vs tool errors) Benefits: - Total code reduction: 99 lines removed - Better maintainability with table-driven tests - Future-proof with official SDK - Consistent with server-side SDK usage - All tests pass (3 test functions, 5 subtests total) Co-Authored-By: Claude Sonnet 4.5 fmt Signed-off-by: Tomasz Janiszewski --- .github/workflows/test.yml | 44 ++++++++ cmd/stackrox-mcp/main.go | 51 +-------- cmd/stackrox-mcp/main_test.go | 12 ++- go.mod | 2 +- integration/fixtures.go | 29 +++++ integration/integration_test.go | 183 ++++++++++++++++++++++++++++++++ internal/app/app.go | 66 ++++++++++++ internal/server/server.go | 18 +++- internal/server/server_test.go | 4 +- internal/testutil/command.go | 18 ++++ internal/testutil/mcp_client.go | 125 ++++++++++++++++++++++ 11 files changed, 498 insertions(+), 54 deletions(-) create mode 100644 integration/fixtures.go create mode 100644 integration/integration_test.go create mode 100644 internal/app/app.go create mode 100644 internal/testutil/command.go create mode 100644 internal/testutil/mcp_client.go diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c8875e5..ee88466 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -45,6 +45,39 @@ jobs: - name: Run E2E smoke test run: make e2e-smoke-test + - name: Set up Java + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: '11' + + - name: Setup proto files + run: ./scripts/setup-proto-files.sh + + - name: Generate proto descriptors + run: make proto-generate + + - name: Download WireMock + run: make mock-download + + - name: Start WireMock + run: make mock-start + + - name: Run integration tests + run: make test-integration-coverage + + - name: Stop WireMock + if: always() + run: make mock-stop + + - name: Upload WireMock logs on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: wiremock-logs + path: wiremock/wiremock.log + if-no-files-found: ignore + - name: Upload test results to Codecov uses: codecov/test-results-action@v1 with: @@ -56,3 +89,14 @@ jobs: files: ./coverage.out token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: false + flags: unit + name: unit-tests + + - name: Upload integration test coverage to Codecov + uses: codecov/codecov-action@v5 + with: + file: ./coverage-integration.out + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: false + flags: integration + name: integration-tests diff --git a/cmd/stackrox-mcp/main.go b/cmd/stackrox-mcp/main.go index c1739df..3810e82 100644 --- a/cmd/stackrox-mcp/main.go +++ b/cmd/stackrox-mcp/main.go @@ -4,28 +4,12 @@ package main import ( "context" "flag" - "log/slog" - "os" - "os/signal" - "syscall" - "github.com/stackrox/stackrox-mcp/internal/client" + "github.com/stackrox/stackrox-mcp/internal/app" "github.com/stackrox/stackrox-mcp/internal/config" "github.com/stackrox/stackrox-mcp/internal/logging" - "github.com/stackrox/stackrox-mcp/internal/server" - "github.com/stackrox/stackrox-mcp/internal/toolsets" - toolsetConfig "github.com/stackrox/stackrox-mcp/internal/toolsets/config" - toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability" ) -// getToolsets initializes and returns all available toolsets. -func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { - return []toolsets.Toolset{ - toolsetConfig.NewToolset(cfg, c), - toolsetVulnerability.NewToolset(cfg, c), - } -} - func main() { logging.SetupLogging() @@ -38,38 +22,9 @@ func main() { logging.Fatal("Failed to load configuration", err) } - // Log full configuration with sensitive data redacted. - slog.Info("Configuration loaded successfully", "config", cfg.Redacted()) - - stackroxClient, err := client.NewClient(&cfg.Central) - if err != nil { - logging.Fatal("Failed to create StackRox client", err) - } - - registry := toolsets.NewRegistry(cfg, getToolsets(cfg, stackroxClient)) - srv := server.NewServer(cfg, registry) - - // Set up context with signal handling for graceful shutdown. - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - err = stackroxClient.Connect(ctx) - if err != nil { - logging.Fatal("Failed to connect to StackRox server", err) - } - - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) - - go func() { - <-sigChan - slog.Info("Received shutdown signal") - cancel() - }() - - slog.Info("Starting StackRox MCP server") + ctx := context.Background() - if err := srv.Start(ctx); err != nil { + if err := app.Run(ctx, cfg, nil, nil); err != nil { logging.Fatal("Server error", err) } } diff --git a/cmd/stackrox-mcp/main_test.go b/cmd/stackrox-mcp/main_test.go index 0202fcb..34bdd04 100644 --- a/cmd/stackrox-mcp/main_test.go +++ b/cmd/stackrox-mcp/main_test.go @@ -14,10 +14,20 @@ import ( "github.com/stackrox/stackrox-mcp/internal/server" "github.com/stackrox/stackrox-mcp/internal/testutil" "github.com/stackrox/stackrox-mcp/internal/toolsets" + toolsetConfig "github.com/stackrox/stackrox-mcp/internal/toolsets/config" + toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +// getToolsets initializes and returns all available toolsets. +func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { + return []toolsets.Toolset{ + toolsetConfig.NewToolset(cfg, c), + toolsetVulnerability.NewToolset(cfg, c), + } +} + func TestGetToolsets(t *testing.T) { allToolsets := getToolsets(&config.Config{}, &client.Client{}) @@ -46,7 +56,7 @@ func TestGracefulShutdown(t *testing.T) { errChan := make(chan error, 1) go func() { - errChan <- srv.Start(ctx) + errChan <- srv.Start(ctx, nil, nil) }() serverURL := "http://" + net.JoinHostPort(cfg.Server.Address, strconv.Itoa(cfg.Server.Port)) diff --git a/go.mod b/go.mod index 486f29d..9673b45 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/stretchr/testify v1.11.1 golang.stackrox.io/grpc-http1 v0.5.1 google.golang.org/grpc v1.79.2 + google.golang.org/protobuf v1.36.10 ) require ( @@ -40,7 +41,6 @@ require ( golang.org/x/text v0.32.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect - google.golang.org/protobuf v1.36.10 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/integration/fixtures.go b/integration/fixtures.go new file mode 100644 index 0000000..fcdb5b9 --- /dev/null +++ b/integration/fixtures.go @@ -0,0 +1,29 @@ +//go:build integration + +package integration + +// Log4ShellFixture contains expected data from log4j_cve.json fixture. +var Log4ShellFixture = struct { + CVEName string + DeploymentCount int + DeploymentNames []string +}{ + CVEName: "CVE-2021-44228", + DeploymentCount: 3, + DeploymentNames: []string{"elasticsearch", "kafka-broker", "spring-boot-app"}, +} + +// AllClustersFixture contains expected data from all_clusters.json fixture. +var AllClustersFixture = struct { + TotalCount int + ClusterNames []string +}{ + TotalCount: 5, + ClusterNames: []string{ + "production-cluster", + "staging-cluster", + "staging-central-cluster", + "development-cluster", + "production-cluster-eu", + }, +} diff --git a/integration/integration_test.go b/integration/integration_test.go new file mode 100644 index 0000000..e0a3f9e --- /dev/null +++ b/integration/integration_test.go @@ -0,0 +1,183 @@ +//go:build integration + +package integration + +import ( + "context" + "io" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stackrox/stackrox-mcp/internal/app" + "github.com/stackrox/stackrox-mcp/internal/config" + "github.com/stackrox/stackrox-mcp/internal/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// setupInitializedClient creates an initialized MCP client for testing. +func setupInitializedClient(t *testing.T) *testutil.MCPTestClient { + t.Helper() + + client, err := createMCPClient(t) + require.NoError(t, err, "Failed to create MCP client") + t.Cleanup(func() { client.Close() }) + + return client +} + +// callToolAndGetResult calls a tool and verifies it succeeds. +func callToolAndGetResult(t *testing.T, client *testutil.MCPTestClient, toolName string, args map[string]any) *mcp.CallToolResult { + t.Helper() + + ctx := context.Background() + result, err := client.CallTool(ctx, toolName, args) + require.NoError(t, err) + testutil.RequireNoError(t, result) + + return result +} + +// getTextContent extracts text from the first content item. +func getTextContent(t *testing.T, result *mcp.CallToolResult) string { + t.Helper() + require.NotEmpty(t, result.Content, "should have content in response") + + textContent, ok := result.Content[0].(*mcp.TextContent) + require.True(t, ok, "expected TextContent, got %T", result.Content[0]) + + return textContent.Text +} + +// TestIntegration_ListTools verifies that all expected tools are registered. +func TestIntegration_ListTools(t *testing.T) { + client := setupInitializedClient(t) + + ctx := context.Background() + result, err := client.ListTools(ctx) + require.NoError(t, err) + + // Verify we have tools registered + assert.NotEmpty(t, result.Tools, "should have tools registered") + + // Check for specific tools we expect + toolNames := make([]string, 0, len(result.Tools)) + for _, tool := range result.Tools { + toolNames = append(toolNames, tool.Name) + } + + assert.Contains(t, toolNames, "get_deployments_for_cve", "should have get_deployments_for_cve tool") + assert.Contains(t, toolNames, "list_clusters", "should have list_clusters tool") +} + +// TestIntegration_ToolCalls tests successful tool calls using table-driven tests. +func TestIntegration_ToolCalls(t *testing.T) { + tests := map[string]struct { + toolName string + args map[string]any + expectedInText []string // strings that must appear in response + }{ + "get_deployments_for_cve with Log4Shell": { + toolName: "get_deployments_for_cve", + args: map[string]any{"cveName": Log4ShellFixture.CVEName}, + expectedInText: Log4ShellFixture.DeploymentNames, + }, + "get_deployments_for_cve with non-existent CVE": { + toolName: "get_deployments_for_cve", + args: map[string]any{"cveName": "CVE-9999-99999"}, + expectedInText: []string{`"deployments":[]`}, + }, + "list_clusters": { + toolName: "list_clusters", + args: map[string]any{}, + expectedInText: AllClustersFixture.ClusterNames, + }, + "get_clusters_with_orchestrator_cve": { + toolName: "get_clusters_with_orchestrator_cve", + args: map[string]any{"cveName": "CVE-2099-00001"}, + expectedInText: []string{`"clusters":`}, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + client := setupInitializedClient(t) + result := callToolAndGetResult(t, client, tt.toolName, tt.args) + + responseText := getTextContent(t, result) + for _, expected := range tt.expectedInText { + assert.Contains(t, responseText, expected) + } + }) + } +} + +// TestIntegration_ToolCallErrors tests error handling using table-driven tests. +func TestIntegration_ToolCallErrors(t *testing.T) { + tests := map[string]struct { + toolName string + args map[string]any + expectedErrorMsg string + }{ + "get_deployments_for_cve missing CVE name": { + toolName: "get_deployments_for_cve", + args: map[string]any{}, + expectedErrorMsg: "cveName", + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + client := setupInitializedClient(t) + + ctx := context.Background() + _, err := client.CallTool(ctx, tt.toolName, tt.args) + + // Validation errors are returned as protocol errors, not tool errors + require.Error(t, err, "should receive protocol error for invalid params") + assert.Contains(t, err.Error(), tt.expectedErrorMsg) + }) + } +} + +// createTestConfig creates a test configuration for the MCP server. +func createTestConfig() *config.Config { + return &config.Config{ + Central: config.CentralConfig{ + URL: "localhost:8081", + AuthType: "static", + APIToken: "test-token-admin", + InsecureSkipTLSVerify: true, + RequestTimeout: 30 * time.Second, + MaxRetries: 3, + InitialBackoff: time.Second, + MaxBackoff: 10 * time.Second, + }, + Server: config.ServerConfig{ + Type: config.ServerTypeStdio, + }, + Tools: config.ToolsConfig{ + Vulnerability: config.ToolsetVulnerabilityConfig{ + Enabled: true, + }, + ConfigManager: config.ToolConfigManagerConfig{ + Enabled: true, + }, + }, + } +} + +// createMCPClient is a helper function that creates an MCP client with the test configuration. +func createMCPClient(t *testing.T) (*testutil.MCPTestClient, error) { + t.Helper() + + cfg := createTestConfig() + + // Create a run function that wraps app.Run with the config + runFunc := func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { + return app.Run(ctx, cfg, stdin, stdout) + } + + return testutil.NewMCPTestClient(t, runFunc) +} diff --git a/internal/app/app.go b/internal/app/app.go new file mode 100644 index 0000000..3aaeba6 --- /dev/null +++ b/internal/app/app.go @@ -0,0 +1,66 @@ +// Package app contains the main application logic for the stackrox-mcp server. +// This is separated from the main package to allow tests to run the server in-process. +package app + +import ( + "context" + "io" + "log/slog" + "os" + "os/signal" + "syscall" + + "github.com/stackrox/stackrox-mcp/internal/client" + "github.com/stackrox/stackrox-mcp/internal/config" + "github.com/stackrox/stackrox-mcp/internal/server" + "github.com/stackrox/stackrox-mcp/internal/toolsets" + toolsetConfig "github.com/stackrox/stackrox-mcp/internal/toolsets/config" + toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability" +) + +// getToolsets initializes and returns all available toolsets. +func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { + return []toolsets.Toolset{ + toolsetConfig.NewToolset(cfg, c), + toolsetVulnerability.NewToolset(cfg, c), + } +} + +// Run executes the MCP server with the given configuration and I/O streams. +// If stdin/stdout are nil, os.Stdin/os.Stdout will be used. +// This function is extracted from main() to allow tests to run the server in-process. +func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io.WriteCloser) error { + // Log full configuration with sensitive data redacted. + slog.Info("Configuration loaded successfully", "config", cfg.Redacted()) + + stackroxClient, err := client.NewClient(&cfg.Central) + if err != nil { + return err + } + + registry := toolsets.NewRegistry(cfg, getToolsets(cfg, stackroxClient)) + srv := server.NewServer(cfg, registry) + + err = stackroxClient.Connect(ctx) + if err != nil { + return err + } + + // Set up signal handling for graceful shutdown. + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + + // Create a cancellable context from the input context + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + go func() { + <-sigChan + slog.Info("Received shutdown signal") + cancel() + }() + + slog.Info("Starting StackRox MCP server") + + return srv.Start(ctx, stdin, stdout) +} diff --git a/internal/server/server.go b/internal/server/server.go index 5309b86..3c2b794 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -3,6 +3,7 @@ package server import ( "context" + "io" "log/slog" "net" "net/http" @@ -50,11 +51,24 @@ func NewServer(cfg *config.Config, registry *toolsets.Registry) *Server { } // Start starts the HTTP server with Streamable HTTP transport. -func (s *Server) Start(ctx context.Context) error { +// If stdin/stdout are provided (non-nil), they will be used for stdio transport. +// If they are nil, os.Stdin/os.Stdout will be used (production mode). +func (s *Server) Start(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { s.registerTools() if s.cfg.Server.Type == config.ServerTypeStdio { - return errors.Wrap(s.mcp.Run(ctx, &mcp.StdioTransport{}), "running mcp over stdio") + var transport mcp.Transport + if stdin != nil && stdout != nil { + // Use custom stdin/stdout (for testing) + transport = &mcp.IOTransport{ + Reader: stdin, + Writer: stdout, + } + } else { + // Use os.Stdin/os.Stdout (production) + transport = &mcp.StdioTransport{} + } + return errors.Wrap(s.mcp.Run(ctx, transport), "running mcp over stdio") } // Create a new ServeMux for routing. diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 156b22a..2968781 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -133,7 +133,7 @@ func TestServer_Start(t *testing.T) { errChan := make(chan error, 1) go func() { - errChan <- srv.Start(ctx) + errChan <- srv.Start(ctx, nil, nil) }() serverURL := "http://" + net.JoinHostPort(cfg.Server.Address, strconv.Itoa(cfg.Server.Port)) @@ -180,7 +180,7 @@ func TestServer_HealthEndpoint(t *testing.T) { errChan := make(chan error, 1) go func() { - errChan <- srv.Start(ctx) + errChan <- srv.Start(ctx, nil, nil) }() serverURL := "http://" + net.JoinHostPort(cfg.Server.Address, strconv.Itoa(cfg.Server.Port)) diff --git a/internal/testutil/command.go b/internal/testutil/command.go new file mode 100644 index 0000000..882a744 --- /dev/null +++ b/internal/testutil/command.go @@ -0,0 +1,18 @@ +package testutil + +import ( + "os/exec" + "strings" +) + +// RunCommand executes a shell command and returns the output and error. +func RunCommand(command string) (string, error) { + parts := strings.Fields(command) + if len(parts) == 0 { + return "", nil + } + + cmd := exec.Command(parts[0], parts[1:]...) + output, err := cmd.CombinedOutput() + return string(output), err +} diff --git a/internal/testutil/mcp_client.go b/internal/testutil/mcp_client.go new file mode 100644 index 0000000..4ec13cd --- /dev/null +++ b/internal/testutil/mcp_client.go @@ -0,0 +1,125 @@ +package testutil + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// MCPTestClient wraps the official MCP SDK client for testing purposes. +type MCPTestClient struct { + session *mcp.ClientSession + cancel context.CancelFunc + errCh chan error + t *testing.T +} + +// ServerRunFunc is a function that runs the MCP server with the given context and I/O streams. +// This allows tests to inject the server run function without creating circular dependencies. +type ServerRunFunc func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error + +// NewMCPTestClient creates a new MCP test client that starts the MCP server in-process with stdio transport. +// The runFunc parameter should be a function that starts the MCP server (typically app.Run wrapped with config). +func NewMCPTestClient(t *testing.T, runFunc ServerRunFunc) (*MCPTestClient, error) { + t.Helper() + + // Create pipes for bidirectional communication + // Server reads from serverStdin, client writes to clientStdout (same pipe) + serverStdin, clientStdout := io.Pipe() + // Server writes to serverStdout, client reads from clientStdin (same pipe) + clientStdin, serverStdout := io.Pipe() + + // Start server in goroutine + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 1) + + go func() { + err := runFunc(ctx, serverStdin, serverStdout) + if err != nil && !errors.Is(err, context.Canceled) { + t.Logf("MCP server error: %v", err) + errCh <- err + } + }() + + // Create MCP client using official SDK + client := mcp.NewClient( + &mcp.Implementation{ + Name: "mcp-test-client", + Version: "1.0.0", + }, + nil, // No custom options needed for basic testing + ) + + // Create IO transport for client + transport := &mcp.IOTransport{ + Reader: clientStdin, // Client reads from this pipe (server writes) + Writer: clientStdout, // Client writes to this pipe (server reads) + } + + // Connect and initialize + session, err := client.Connect(ctx, transport, &mcp.ClientSessionOptions{}) + if err != nil { + cancel() + clientStdout.Close() + clientStdin.Close() + return nil, err + } + + return &MCPTestClient{ + session: session, + cancel: cancel, + errCh: errCh, + t: t, + }, nil +} + +// Close stops the MCP server and cleans up resources. +func (c *MCPTestClient) Close() error { + if err := c.session.Close(); err != nil { + c.t.Logf("Error closing session: %v", err) + } + c.cancel() + + // Wait for server to finish (with timeout) + select { + case err := <-c.errCh: + if err != nil && !errors.Is(err, context.Canceled) { + return err + } + default: + // Server is still running or finished cleanly + } + + return nil +} + +// ListTools returns all available tools from the server. +func (c *MCPTestClient) ListTools(ctx context.Context) (*mcp.ListToolsResult, error) { + return c.session.ListTools(ctx, nil) +} + +// CallTool invokes a tool with the given name and arguments. +func (c *MCPTestClient) CallTool(ctx context.Context, toolName string, args map[string]any) (*mcp.CallToolResult, error) { + return c.session.CallTool(ctx, &mcp.CallToolParams{ + Name: toolName, + Arguments: args, + }) +} + +// RequireNoError asserts that the tool call result does not contain an error. +func RequireNoError(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + if result.IsError { + // Extract error message from content + errMsg := "unknown error" + if len(result.Content) > 0 { + if textContent, ok := result.Content[0].(*mcp.TextContent); ok { + errMsg = textContent.Text + } + } + t.Fatalf("expected no error, got: %s", errMsg) + } +} From 3c617df6f4d3f5e6b3954d78da76f444af3b7614 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Fri, 6 Mar 2026 15:45:41 +0100 Subject: [PATCH 02/19] ci: fix golangci-lint style issues Fixed 17 linter issues across multiple packages: - errcheck: Added error checks for Close() calls - gosec: Added security comment for test utility - lll: Split long function signature - nlreturn: Added blank lines before returns - wrapcheck: Wrapped external package errors with context - wsl: Fixed whitespace requirements - noctx: Used CommandContext instead of Command Co-Authored-By: Claude Sonnet 4.5 --- internal/app/app.go | 7 ++++--- internal/server/server.go | 1 + internal/testutil/command.go | 5 ++++- internal/testutil/mcp_client.go | 33 +++++++++++++++++++++++++++------ 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index 3aaeba6..567cc0f 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -10,6 +10,7 @@ import ( "os/signal" "syscall" + "github.com/pkg/errors" "github.com/stackrox/stackrox-mcp/internal/client" "github.com/stackrox/stackrox-mcp/internal/config" "github.com/stackrox/stackrox-mcp/internal/server" @@ -35,7 +36,7 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io stackroxClient, err := client.NewClient(&cfg.Central) if err != nil { - return err + return errors.Wrap(err, "failed to create client") } registry := toolsets.NewRegistry(cfg, getToolsets(cfg, stackroxClient)) @@ -43,7 +44,7 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io err = stackroxClient.Connect(ctx) if err != nil { - return err + return errors.Wrap(err, "failed to connect to central") } // Set up signal handling for graceful shutdown. @@ -62,5 +63,5 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io slog.Info("Starting StackRox MCP server") - return srv.Start(ctx, stdin, stdout) + return errors.Wrap(srv.Start(ctx, stdin, stdout), "failed to start server") } diff --git a/internal/server/server.go b/internal/server/server.go index 3c2b794..6e6c571 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -68,6 +68,7 @@ func (s *Server) Start(ctx context.Context, stdin io.ReadCloser, stdout io.Write // Use os.Stdin/os.Stdout (production) transport = &mcp.StdioTransport{} } + return errors.Wrap(s.mcp.Run(ctx, transport), "running mcp over stdio") } diff --git a/internal/testutil/command.go b/internal/testutil/command.go index 882a744..45dd6b4 100644 --- a/internal/testutil/command.go +++ b/internal/testutil/command.go @@ -1,6 +1,7 @@ package testutil import ( + "context" "os/exec" "strings" ) @@ -12,7 +13,9 @@ func RunCommand(command string) (string, error) { return "", nil } - cmd := exec.Command(parts[0], parts[1:]...) + // #nosec G204 - This is a test utility function with controlled input + cmd := exec.CommandContext(context.Background(), parts[0], parts[1:]...) output, err := cmd.CombinedOutput() + return string(output), err } diff --git a/internal/testutil/mcp_client.go b/internal/testutil/mcp_client.go index 4ec13cd..5f7f828 100644 --- a/internal/testutil/mcp_client.go +++ b/internal/testutil/mcp_client.go @@ -40,6 +40,7 @@ func NewMCPTestClient(t *testing.T, runFunc ServerRunFunc) (*MCPTestClient, erro err := runFunc(ctx, serverStdin, serverStdout) if err != nil && !errors.Is(err, context.Canceled) { t.Logf("MCP server error: %v", err) + errCh <- err } }() @@ -63,9 +64,11 @@ func NewMCPTestClient(t *testing.T, runFunc ServerRunFunc) (*MCPTestClient, erro session, err := client.Connect(ctx, transport, &mcp.ClientSessionOptions{}) if err != nil { cancel() - clientStdout.Close() - clientStdin.Close() - return nil, err + + _ = clientStdout.Close() + _ = clientStdin.Close() + + return nil, errors.New("failed to connect to MCP server: " + err.Error()) } return &MCPTestClient{ @@ -81,6 +84,7 @@ func (c *MCPTestClient) Close() error { if err := c.session.Close(); err != nil { c.t.Logf("Error closing session: %v", err) } + c.cancel() // Wait for server to finish (with timeout) @@ -98,28 +102,45 @@ func (c *MCPTestClient) Close() error { // ListTools returns all available tools from the server. func (c *MCPTestClient) ListTools(ctx context.Context) (*mcp.ListToolsResult, error) { - return c.session.ListTools(ctx, nil) + result, err := c.session.ListTools(ctx, nil) + if err != nil { + return nil, errors.New("failed to list tools: " + err.Error()) + } + + return result, nil } // CallTool invokes a tool with the given name and arguments. -func (c *MCPTestClient) CallTool(ctx context.Context, toolName string, args map[string]any) (*mcp.CallToolResult, error) { - return c.session.CallTool(ctx, &mcp.CallToolParams{ +func (c *MCPTestClient) CallTool( + ctx context.Context, + toolName string, + args map[string]any, +) (*mcp.CallToolResult, error) { + result, err := c.session.CallTool(ctx, &mcp.CallToolParams{ Name: toolName, Arguments: args, }) + if err != nil { + return nil, errors.New("failed to call tool: " + err.Error()) + } + + return result, nil } // RequireNoError asserts that the tool call result does not contain an error. func RequireNoError(t *testing.T, result *mcp.CallToolResult) { t.Helper() + if result.IsError { // Extract error message from content errMsg := "unknown error" + if len(result.Content) > 0 { if textContent, ok := result.Content[0].(*mcp.TextContent); ok { errMsg = textContent.Text } } + t.Fatalf("expected no error, got: %s", errMsg) } } From b3fcf5cfbedbad15f7ab30324778525eb65397d2 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Fri, 6 Mar 2026 17:31:27 +0100 Subject: [PATCH 03/19] refactor: extract helper methods to fix funlen linter issue Split the Start function into smaller helper methods: - startStdio: handles stdio transport setup - startHTTP: handles HTTP server lifecycle This reduces the Start function from 61 to 11 lines, resolving the funlen linter error (max 60 lines). Co-Authored-By: Claude Sonnet 4.5 --- internal/server/server.go | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/internal/server/server.go b/internal/server/server.go index 6e6c571..8740f5f 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -57,21 +57,29 @@ func (s *Server) Start(ctx context.Context, stdin io.ReadCloser, stdout io.Write s.registerTools() if s.cfg.Server.Type == config.ServerTypeStdio { - var transport mcp.Transport - if stdin != nil && stdout != nil { - // Use custom stdin/stdout (for testing) - transport = &mcp.IOTransport{ - Reader: stdin, - Writer: stdout, - } - } else { - // Use os.Stdin/os.Stdout (production) - transport = &mcp.StdioTransport{} - } + return s.startStdio(ctx, stdin, stdout) + } - return errors.Wrap(s.mcp.Run(ctx, transport), "running mcp over stdio") + return s.startHTTP(ctx) +} + +func (s *Server) startStdio(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { + var transport mcp.Transport + if stdin != nil && stdout != nil { + // Use custom stdin/stdout (for testing) + transport = &mcp.IOTransport{ + Reader: stdin, + Writer: stdout, + } + } else { + // Use os.Stdin/os.Stdout (production) + transport = &mcp.StdioTransport{} } + return errors.Wrap(s.mcp.Run(ctx, transport), "running mcp over stdio") +} + +func (s *Server) startHTTP(ctx context.Context) error { // Create a new ServeMux for routing. mux := http.NewServeMux() s.registerRouteHealth(mux) From 630b3daaad0b2006f79b40537904fe1bb38f7825 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Mon, 9 Mar 2026 12:08:32 +0100 Subject: [PATCH 04/19] fix Signed-off-by: Tomasz Janiszewski --- .github/workflows/test.yml | 2 +- go.mod | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ee88466..c885c5b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -95,7 +95,7 @@ jobs: - name: Upload integration test coverage to Codecov uses: codecov/codecov-action@v5 with: - file: ./coverage-integration.out + files: ./coverage-integration.out token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: false flags: integration diff --git a/go.mod b/go.mod index 9673b45..486f29d 100644 --- a/go.mod +++ b/go.mod @@ -11,7 +11,6 @@ require ( github.com/stretchr/testify v1.11.1 golang.stackrox.io/grpc-http1 v0.5.1 google.golang.org/grpc v1.79.2 - google.golang.org/protobuf v1.36.10 ) require ( @@ -41,6 +40,7 @@ require ( golang.org/x/text v0.32.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect + google.golang.org/protobuf v1.36.10 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) From e911b97755601c2bf9e7ed17b25765e28e3a2a52 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Mon, 9 Mar 2026 15:51:05 +0100 Subject: [PATCH 05/19] fix Signed-off-by: Tomasz Janiszewski --- Makefile | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Makefile b/Makefile index b74dcfa..8a727a7 100644 --- a/Makefile +++ b/Makefile @@ -70,6 +70,10 @@ test-coverage-and-junit: ## Run unit tests with coverage and junit output go install github.com/jstemmer/go-junit-report/v2@v2.1.0 $(GOTEST) -v -cover -race -coverprofile=$(COVERAGE_OUT) ./... 2>&1 | go-junit-report -set-exit-code -iocopy -out $(JUNIT_OUT) +.PHONY: test-integration-coverage +test-integration-coverage: ## Run integration tests with coverage + $(GOTEST) -v -cover -race -tags=integration -coverprofile=coverage-integration.out ./integration + .PHONY: coverage-html coverage-html: test ## Generate and open HTML coverage report $(GOCMD) tool cover -html=$(COVERAGE_OUT) From 8550a0c4d0ee86c277e2c91dc366d76fc655edb8 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Mon, 9 Mar 2026 18:01:58 +0100 Subject: [PATCH 06/19] fix Signed-off-by: Tomasz Janiszewski --- internal/app/app_test.go | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 internal/app/app_test.go diff --git a/internal/app/app_test.go b/internal/app/app_test.go new file mode 100644 index 0000000..19a2111 --- /dev/null +++ b/internal/app/app_test.go @@ -0,0 +1,33 @@ +package app + +import ( + "testing" + + "github.com/stackrox/stackrox-mcp/internal/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetToolsets(t *testing.T) { + cfg := &config.Config{ + Central: config.CentralConfig{ + URL: "https://example.com", + }, + Tools: config.ToolsConfig{ + ConfigManager: config.ToolConfigManagerConfig{ + Enabled: false, + }, + Vulnerability: config.ToolsetVulnerabilityConfig{ + Enabled: true, + }, + }, + } + + // We can't create a real client without a valid connection, + // so we pass nil and just test that getToolsets returns something + toolsets := getToolsets(cfg, nil) + + require.NotNil(t, toolsets, "getToolsets should return a non-nil slice") + assert.NotEmpty(t, toolsets, "getToolsets should return at least one toolset") + assert.Len(t, toolsets, 2, "getToolsets should return 2 toolsets") +} From 07f28026f5159d66910a0435586ddcaf4aa8b08e Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:47:48 +0100 Subject: [PATCH 07/19] fix: export GetToolsets function for reusability Change getToolsets to GetToolsets to make it available for use in test files outside the app package. Addresses PR #50 review comment #1 Co-Authored-By: Claude Sonnet 4.5 --- internal/app/app.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index 567cc0f..7006e54 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -19,8 +19,8 @@ import ( toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability" ) -// getToolsets initializes and returns all available toolsets. -func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { +// GetToolsets initializes and returns all available toolsets. +func GetToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { return []toolsets.Toolset{ toolsetConfig.NewToolset(cfg, c), toolsetVulnerability.NewToolset(cfg, c), @@ -39,7 +39,7 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io return errors.Wrap(err, "failed to create client") } - registry := toolsets.NewRegistry(cfg, getToolsets(cfg, stackroxClient)) + registry := toolsets.NewRegistry(cfg, GetToolsets(cfg, stackroxClient)) srv := server.NewServer(cfg, registry) err = stackroxClient.Connect(ctx) From 6424c85cc8ad4a58c9e54a7ca271a4f7552c963d Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:48:15 +0100 Subject: [PATCH 08/19] fix: remove duplicate getToolsets and use app.GetToolsets Remove duplicate getToolsets function from main_test.go and use the exported app.GetToolsets instead. Also remove the duplicate TestGetToolsets test which will be enhanced in app_test.go. Addresses PR #50 review comment #2 Co-Authored-By: Claude Sonnet 4.5 --- cmd/stackrox-mcp/main_test.go | 26 ++------------------------ 1 file changed, 2 insertions(+), 24 deletions(-) diff --git a/cmd/stackrox-mcp/main_test.go b/cmd/stackrox-mcp/main_test.go index 34bdd04..3222127 100644 --- a/cmd/stackrox-mcp/main_test.go +++ b/cmd/stackrox-mcp/main_test.go @@ -9,37 +9,15 @@ import ( "testing" "time" + "github.com/stackrox/stackrox-mcp/internal/app" "github.com/stackrox/stackrox-mcp/internal/client" "github.com/stackrox/stackrox-mcp/internal/config" "github.com/stackrox/stackrox-mcp/internal/server" "github.com/stackrox/stackrox-mcp/internal/testutil" "github.com/stackrox/stackrox-mcp/internal/toolsets" - toolsetConfig "github.com/stackrox/stackrox-mcp/internal/toolsets/config" - toolsetVulnerability "github.com/stackrox/stackrox-mcp/internal/toolsets/vulnerability" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// getToolsets initializes and returns all available toolsets. -func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { - return []toolsets.Toolset{ - toolsetConfig.NewToolset(cfg, c), - toolsetVulnerability.NewToolset(cfg, c), - } -} - -func TestGetToolsets(t *testing.T) { - allToolsets := getToolsets(&config.Config{}, &client.Client{}) - - toolsetNames := []string{} - for _, toolset := range allToolsets { - toolsetNames = append(toolsetNames, toolset.GetName()) - } - - assert.Contains(t, toolsetNames, "config_manager") - assert.Contains(t, toolsetNames, "vulnerability") -} - func TestGracefulShutdown(t *testing.T) { // Set up minimal valid config. config.LoadConfig() validates configuration. t.Setenv("STACKROX_MCP__TOOLS__VULNERABILITY__ENABLED", "true") @@ -49,7 +27,7 @@ func TestGracefulShutdown(t *testing.T) { require.NotNil(t, cfg) cfg.Server.Port = testutil.GetPortForTest(t) - registry := toolsets.NewRegistry(cfg, getToolsets(cfg, &client.Client{})) + registry := toolsets.NewRegistry(cfg, app.GetToolsets(cfg, &client.Client{})) srv := server.NewServer(cfg, registry) ctx, cancel := context.WithCancel(context.Background()) From c7273238d8777861f47c36d627f57aebde458556 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:48:33 +0100 Subject: [PATCH 09/19] test: move comprehensive TestGetToolsets to app package Replace simple test with comprehensive version that verifies specific toolset names are present. This test was previously in cmd/stackrox-mcp/main_test.go and is now in the proper location alongside the GetToolsets function. Addresses PR #50 review comment #3 Co-Authored-By: Claude Sonnet 4.5 --- internal/app/app_test.go | 28 ++++++++-------------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/internal/app/app_test.go b/internal/app/app_test.go index 19a2111..8e1c742 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -3,31 +3,19 @@ package app import ( "testing" + "github.com/stackrox/stackrox-mcp/internal/client" "github.com/stackrox/stackrox-mcp/internal/config" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) func TestGetToolsets(t *testing.T) { - cfg := &config.Config{ - Central: config.CentralConfig{ - URL: "https://example.com", - }, - Tools: config.ToolsConfig{ - ConfigManager: config.ToolConfigManagerConfig{ - Enabled: false, - }, - Vulnerability: config.ToolsetVulnerabilityConfig{ - Enabled: true, - }, - }, - } + allToolsets := GetToolsets(&config.Config{}, &client.Client{}) - // We can't create a real client without a valid connection, - // so we pass nil and just test that getToolsets returns something - toolsets := getToolsets(cfg, nil) + toolsetNames := []string{} + for _, toolset := range allToolsets { + toolsetNames = append(toolsetNames, toolset.GetName()) + } - require.NotNil(t, toolsets, "getToolsets should return a non-nil slice") - assert.NotEmpty(t, toolsets, "getToolsets should return at least one toolset") - assert.Len(t, toolsets, 2, "getToolsets should return 2 toolsets") + assert.Contains(t, toolsetNames, "config_manager") + assert.Contains(t, toolsetNames, "vulnerability") } From f3c3f8ef9cc78a04eef8ef29d294eccf9fcdb159 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:48:44 +0100 Subject: [PATCH 10/19] docs: remove irrelevant stdin/stdout comment from app.go Remove implementation detail comment about stdin/stdout behavior that clutters the high-level API documentation. Addresses PR #50 review comment #4 Co-Authored-By: Claude Sonnet 4.5 --- internal/app/app.go | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/app/app.go b/internal/app/app.go index 7006e54..bf4f724 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -28,7 +28,6 @@ func GetToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset { } // Run executes the MCP server with the given configuration and I/O streams. -// If stdin/stdout are nil, os.Stdin/os.Stdout will be used. // This function is extracted from main() to allow tests to run the server in-process. func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io.WriteCloser) error { // Log full configuration with sensitive data redacted. From 4aadc98a47323c06c0aa05a0710b72e6ade9c815 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:48:54 +0100 Subject: [PATCH 11/19] docs: remove '(production mode)' annotation from server.go Remove parenthetical note that adds unnecessary context noise to the documentation. Addresses PR #50 review comment #6 Co-Authored-By: Claude Sonnet 4.5 --- internal/server/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/server/server.go b/internal/server/server.go index 8740f5f..c8d4e9c 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -52,7 +52,7 @@ func NewServer(cfg *config.Config, registry *toolsets.Registry) *Server { // Start starts the HTTP server with Streamable HTTP transport. // If stdin/stdout are provided (non-nil), they will be used for stdio transport. -// If they are nil, os.Stdin/os.Stdout will be used (production mode). +// If they are nil, os.Stdin/os.Stdout will be used. func (s *Server) Start(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { s.registerTools() From 8e732257141c326f64a59ba9afd35ec8a11c6343 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:49:03 +0100 Subject: [PATCH 12/19] docs: remove '(for testing)' annotation from server.go Remove parenthetical note that adds unnecessary context noise to the documentation. Addresses PR #50 review comment #7 Co-Authored-By: Claude Sonnet 4.5 --- internal/server/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/server/server.go b/internal/server/server.go index c8d4e9c..2307820 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -66,7 +66,7 @@ func (s *Server) Start(ctx context.Context, stdin io.ReadCloser, stdout io.Write func (s *Server) startStdio(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { var transport mcp.Transport if stdin != nil && stdout != nil { - // Use custom stdin/stdout (for testing) + // Use custom stdin/stdout transport = &mcp.IOTransport{ Reader: stdin, Writer: stdout, From 46bb0fbbea9400a93bc413f0eef8b8731c924f51 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:49:11 +0100 Subject: [PATCH 13/19] docs: remove '(production)' annotation from server.go Remove parenthetical note that adds unnecessary context noise to the documentation. Addresses PR #50 review comment #8 Co-Authored-By: Claude Sonnet 4.5 --- internal/server/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/server/server.go b/internal/server/server.go index 2307820..fc557ae 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -72,7 +72,7 @@ func (s *Server) startStdio(ctx context.Context, stdin io.ReadCloser, stdout io. Writer: stdout, } } else { - // Use os.Stdin/os.Stdout (production) + // Use os.Stdin/os.Stdout transport = &mcp.StdioTransport{} } From 25d1d91c24573807f9eed508df9496960d2aea1c Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:49:24 +0100 Subject: [PATCH 14/19] docs: add full paths to fixture comments Add complete file paths to make it easier to locate the actual fixture files without searching. Includes paths for both Log4ShellFixture and AllClustersFixture. Addresses PR #50 review comment #9 Co-Authored-By: Claude Sonnet 4.5 --- integration/fixtures.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integration/fixtures.go b/integration/fixtures.go index fcdb5b9..d7bbc5a 100644 --- a/integration/fixtures.go +++ b/integration/fixtures.go @@ -2,7 +2,7 @@ package integration -// Log4ShellFixture contains expected data from log4j_cve.json fixture. +// Log4ShellFixture contains expected data from wiremock/fixtures/deployments/log4j_cve.json fixture. var Log4ShellFixture = struct { CVEName string DeploymentCount int @@ -13,7 +13,7 @@ var Log4ShellFixture = struct { DeploymentNames: []string{"elasticsearch", "kafka-broker", "spring-boot-app"}, } -// AllClustersFixture contains expected data from all_clusters.json fixture. +// AllClustersFixture contains expected data from wiremock/fixtures/clusters/all_clusters.json fixture. var AllClustersFixture = struct { TotalCount int ClusterNames []string From c79028d221080e8af6ed51e7e7a6c686fbfd7e07 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:50:04 +0100 Subject: [PATCH 15/19] refactor: move helper functions to top of integration_test.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reorganize helper functions to the top of the file after imports. New order: imports → helpers (createTestConfig, createMCPClient, setupInitializedClient, callToolAndGetResult, getTextContent) → tests. This improves code organization and readability. Addresses PR #50 review comment #10 Co-Authored-By: Claude Sonnet 4.5 --- integration/integration_test.go | 82 ++++++++++++++++----------------- 1 file changed, 41 insertions(+), 41 deletions(-) diff --git a/integration/integration_test.go b/integration/integration_test.go index e0a3f9e..58b1681 100644 --- a/integration/integration_test.go +++ b/integration/integration_test.go @@ -16,6 +16,47 @@ import ( "github.com/stretchr/testify/require" ) +// createTestConfig creates a test configuration for the MCP server. +func createTestConfig() *config.Config { + return &config.Config{ + Central: config.CentralConfig{ + URL: "localhost:8081", + AuthType: "static", + APIToken: "test-token-admin", + InsecureSkipTLSVerify: true, + RequestTimeout: 30 * time.Second, + MaxRetries: 3, + InitialBackoff: time.Second, + MaxBackoff: 10 * time.Second, + }, + Server: config.ServerConfig{ + Type: config.ServerTypeStdio, + }, + Tools: config.ToolsConfig{ + Vulnerability: config.ToolsetVulnerabilityConfig{ + Enabled: true, + }, + ConfigManager: config.ToolConfigManagerConfig{ + Enabled: true, + }, + }, + } +} + +// createMCPClient is a helper function that creates an MCP client with the test configuration. +func createMCPClient(t *testing.T) (*testutil.MCPTestClient, error) { + t.Helper() + + cfg := createTestConfig() + + // Create a run function that wraps app.Run with the config + runFunc := func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { + return app.Run(ctx, cfg, stdin, stdout) + } + + return testutil.NewMCPTestClient(t, runFunc) +} + // setupInitializedClient creates an initialized MCP client for testing. func setupInitializedClient(t *testing.T) *testutil.MCPTestClient { t.Helper() @@ -140,44 +181,3 @@ func TestIntegration_ToolCallErrors(t *testing.T) { }) } } - -// createTestConfig creates a test configuration for the MCP server. -func createTestConfig() *config.Config { - return &config.Config{ - Central: config.CentralConfig{ - URL: "localhost:8081", - AuthType: "static", - APIToken: "test-token-admin", - InsecureSkipTLSVerify: true, - RequestTimeout: 30 * time.Second, - MaxRetries: 3, - InitialBackoff: time.Second, - MaxBackoff: 10 * time.Second, - }, - Server: config.ServerConfig{ - Type: config.ServerTypeStdio, - }, - Tools: config.ToolsConfig{ - Vulnerability: config.ToolsetVulnerabilityConfig{ - Enabled: true, - }, - ConfigManager: config.ToolConfigManagerConfig{ - Enabled: true, - }, - }, - } -} - -// createMCPClient is a helper function that creates an MCP client with the test configuration. -func createMCPClient(t *testing.T) (*testutil.MCPTestClient, error) { - t.Helper() - - cfg := createTestConfig() - - // Create a run function that wraps app.Run with the config - runFunc := func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { - return app.Run(ctx, cfg, stdin, stdout) - } - - return testutil.NewMCPTestClient(t, runFunc) -} From 8b90a9bdff42b8a598885a3e8513f8ae96054183 Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:50:18 +0100 Subject: [PATCH 16/19] refactor: simplify RequireNoError with early return pattern Reduce nesting level by using early return for the success case. This improves readability by handling the happy path first. Addresses PR #50 review comment #11 Co-Authored-By: Claude Sonnet 4.5 --- internal/testutil/mcp_client.go | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/internal/testutil/mcp_client.go b/internal/testutil/mcp_client.go index 5f7f828..2e78f70 100644 --- a/internal/testutil/mcp_client.go +++ b/internal/testutil/mcp_client.go @@ -131,16 +131,18 @@ func (c *MCPTestClient) CallTool( func RequireNoError(t *testing.T, result *mcp.CallToolResult) { t.Helper() - if result.IsError { - // Extract error message from content - errMsg := "unknown error" - - if len(result.Content) > 0 { - if textContent, ok := result.Content[0].(*mcp.TextContent); ok { - errMsg = textContent.Text - } - } + if !result.IsError { + return + } + + // Extract error message from content + errMsg := "unknown error" - t.Fatalf("expected no error, got: %s", errMsg) + if len(result.Content) > 0 { + if textContent, ok := result.Content[0].(*mcp.TextContent); ok { + errMsg = textContent.Text + } } + + t.Fatalf("expected no error, got: %s", errMsg) } From afc96aa57d1f2807fff87aa2a356d55e33cc200e Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:50:28 +0100 Subject: [PATCH 17/19] chore: delete unused RunCommand function Remove internal/testutil/command.go which contains the unused RunCommand function. Verified by searching the codebase - no actual usage found in any test or production code. Addresses PR #50 review comment #12 Co-Authored-By: Claude Sonnet 4.5 --- internal/testutil/command.go | 21 --------------------- 1 file changed, 21 deletions(-) delete mode 100644 internal/testutil/command.go diff --git a/internal/testutil/command.go b/internal/testutil/command.go deleted file mode 100644 index 45dd6b4..0000000 --- a/internal/testutil/command.go +++ /dev/null @@ -1,21 +0,0 @@ -package testutil - -import ( - "context" - "os/exec" - "strings" -) - -// RunCommand executes a shell command and returns the output and error. -func RunCommand(command string) (string, error) { - parts := strings.Fields(command) - if len(parts) == 0 { - return "", nil - } - - // #nosec G204 - This is a test utility function with controlled input - cmd := exec.CommandContext(context.Background(), parts[0], parts[1:]...) - output, err := cmd.CombinedOutput() - - return string(output), err -} From 6708e6b53c1ec3f7ec67b8d726d4d84c1580bf9d Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:55:58 +0100 Subject: [PATCH 18/19] fix: create cancellable context before client connection Move cancellable context creation to the beginning of Run() so that the entire server lifecycle, including the client connection, can be properly cancelled. This allows graceful shutdown to interrupt the connection phase if needed. The signal handler is also set up earlier to ensure interrupts are caught as soon as possible. Addresses PR #50 review comment #5 Co-Authored-By: Claude Sonnet 4.5 --- internal/app/app.go | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index bf4f724..22f6b27 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -33,6 +33,20 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io // Log full configuration with sensitive data redacted. slog.Info("Configuration loaded successfully", "config", cfg.Redacted()) + // Create a cancellable context for the entire server lifecycle + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + // Set up signal handling for graceful shutdown. + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + + go func() { + <-sigChan + slog.Info("Received shutdown signal") + cancel() + }() + stackroxClient, err := client.NewClient(&cfg.Central) if err != nil { return errors.Wrap(err, "failed to create client") @@ -46,20 +60,6 @@ func Run(ctx context.Context, cfg *config.Config, stdin io.ReadCloser, stdout io return errors.Wrap(err, "failed to connect to central") } - // Set up signal handling for graceful shutdown. - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) - - // Create a cancellable context from the input context - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - go func() { - <-sigChan - slog.Info("Received shutdown signal") - cancel() - }() - slog.Info("Starting StackRox MCP server") return errors.Wrap(srv.Start(ctx, stdin, stdout), "failed to start server") From 46d77093d6ef6a28d8fa7f98bb49a352506b41ea Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 17 Mar 2026 13:57:40 +0100 Subject: [PATCH 19/19] refactor: consolidate integration test helpers in testutil Move all integration test helper functions from integration package to internal/testutil/integration_helpers.go to avoid duplication and centralize test utilities. This improves code organization and makes helpers reusable across different test packages. Helpers moved: - createTestConfig -> CreateIntegrationTestConfig - createMCPClient -> CreateIntegrationMCPClient - setupInitializedClient -> SetupInitializedClient (made generic) - callToolAndGetResult -> CallToolAndGetResult - getTextContent -> GetTextContent The SetupInitializedClient is now generic and accepts a client creation function, making it reusable for different test scenarios. Addresses PR #50 review comment #13 Co-Authored-By: Claude Sonnet 4.5 --- integration/integration_test.go | 90 ++--------------------- internal/testutil/integration_helpers.go | 91 ++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 85 deletions(-) create mode 100644 internal/testutil/integration_helpers.go diff --git a/integration/integration_test.go b/integration/integration_test.go index 58b1681..1addc22 100644 --- a/integration/integration_test.go +++ b/integration/integration_test.go @@ -4,96 +4,16 @@ package integration import ( "context" - "io" "testing" - "time" - "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/stackrox/stackrox-mcp/internal/app" - "github.com/stackrox/stackrox-mcp/internal/config" "github.com/stackrox/stackrox-mcp/internal/testutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// createTestConfig creates a test configuration for the MCP server. -func createTestConfig() *config.Config { - return &config.Config{ - Central: config.CentralConfig{ - URL: "localhost:8081", - AuthType: "static", - APIToken: "test-token-admin", - InsecureSkipTLSVerify: true, - RequestTimeout: 30 * time.Second, - MaxRetries: 3, - InitialBackoff: time.Second, - MaxBackoff: 10 * time.Second, - }, - Server: config.ServerConfig{ - Type: config.ServerTypeStdio, - }, - Tools: config.ToolsConfig{ - Vulnerability: config.ToolsetVulnerabilityConfig{ - Enabled: true, - }, - ConfigManager: config.ToolConfigManagerConfig{ - Enabled: true, - }, - }, - } -} - -// createMCPClient is a helper function that creates an MCP client with the test configuration. -func createMCPClient(t *testing.T) (*testutil.MCPTestClient, error) { - t.Helper() - - cfg := createTestConfig() - - // Create a run function that wraps app.Run with the config - runFunc := func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { - return app.Run(ctx, cfg, stdin, stdout) - } - - return testutil.NewMCPTestClient(t, runFunc) -} - -// setupInitializedClient creates an initialized MCP client for testing. -func setupInitializedClient(t *testing.T) *testutil.MCPTestClient { - t.Helper() - - client, err := createMCPClient(t) - require.NoError(t, err, "Failed to create MCP client") - t.Cleanup(func() { client.Close() }) - - return client -} - -// callToolAndGetResult calls a tool and verifies it succeeds. -func callToolAndGetResult(t *testing.T, client *testutil.MCPTestClient, toolName string, args map[string]any) *mcp.CallToolResult { - t.Helper() - - ctx := context.Background() - result, err := client.CallTool(ctx, toolName, args) - require.NoError(t, err) - testutil.RequireNoError(t, result) - - return result -} - -// getTextContent extracts text from the first content item. -func getTextContent(t *testing.T, result *mcp.CallToolResult) string { - t.Helper() - require.NotEmpty(t, result.Content, "should have content in response") - - textContent, ok := result.Content[0].(*mcp.TextContent) - require.True(t, ok, "expected TextContent, got %T", result.Content[0]) - - return textContent.Text -} - // TestIntegration_ListTools verifies that all expected tools are registered. func TestIntegration_ListTools(t *testing.T) { - client := setupInitializedClient(t) + client := testutil.SetupInitializedClient(t, testutil.CreateIntegrationMCPClient) ctx := context.Background() result, err := client.ListTools(ctx) @@ -143,10 +63,10 @@ func TestIntegration_ToolCalls(t *testing.T) { for name, tt := range tests { t.Run(name, func(t *testing.T) { - client := setupInitializedClient(t) - result := callToolAndGetResult(t, client, tt.toolName, tt.args) + client := testutil.SetupInitializedClient(t, testutil.CreateIntegrationMCPClient) + result := testutil.CallToolAndGetResult(t, client, tt.toolName, tt.args) - responseText := getTextContent(t, result) + responseText := testutil.GetTextContent(t, result) for _, expected := range tt.expectedInText { assert.Contains(t, responseText, expected) } @@ -170,7 +90,7 @@ func TestIntegration_ToolCallErrors(t *testing.T) { for name, tt := range tests { t.Run(name, func(t *testing.T) { - client := setupInitializedClient(t) + client := testutil.SetupInitializedClient(t, testutil.CreateIntegrationMCPClient) ctx := context.Background() _, err := client.CallTool(ctx, tt.toolName, tt.args) diff --git a/internal/testutil/integration_helpers.go b/internal/testutil/integration_helpers.go new file mode 100644 index 0000000..21b48e7 --- /dev/null +++ b/internal/testutil/integration_helpers.go @@ -0,0 +1,91 @@ +//go:build integration + +package testutil + +import ( + "context" + "io" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stackrox/stackrox-mcp/internal/app" + "github.com/stackrox/stackrox-mcp/internal/config" + "github.com/stretchr/testify/require" +) + +// CreateIntegrationTestConfig creates a test configuration for integration tests. +// This connects to a local WireMock instance at localhost:8081. +func CreateIntegrationTestConfig() *config.Config { + return &config.Config{ + Central: config.CentralConfig{ + URL: "localhost:8081", + AuthType: "static", + APIToken: "test-token-admin", + InsecureSkipTLSVerify: true, + RequestTimeout: 30 * time.Second, + MaxRetries: 3, + InitialBackoff: time.Second, + MaxBackoff: 10 * time.Second, + }, + Server: config.ServerConfig{ + Type: config.ServerTypeStdio, + }, + Tools: config.ToolsConfig{ + Vulnerability: config.ToolsetVulnerabilityConfig{ + Enabled: true, + }, + ConfigManager: config.ToolConfigManagerConfig{ + Enabled: true, + }, + }, + } +} + +// CreateIntegrationMCPClient creates an MCP client with integration test configuration. +func CreateIntegrationMCPClient(t *testing.T) (*MCPTestClient, error) { + t.Helper() + + cfg := CreateIntegrationTestConfig() + + // Create a run function that wraps app.Run with the config + runFunc := func(ctx context.Context, stdin io.ReadCloser, stdout io.WriteCloser) error { + return app.Run(ctx, cfg, stdin, stdout) + } + + return NewMCPTestClient(t, runFunc) +} + +// SetupInitializedClient creates an initialized MCP client for testing with automatic cleanup. +func SetupInitializedClient(t *testing.T, createClient func(*testing.T) (*MCPTestClient, error)) *MCPTestClient { + t.Helper() + + client, err := createClient(t) + require.NoError(t, err, "Failed to create MCP client") + t.Cleanup(func() { client.Close() }) + + return client +} + +// CallToolAndGetResult calls a tool and verifies it succeeds. +func CallToolAndGetResult(t *testing.T, client *MCPTestClient, toolName string, args map[string]any) *mcp.CallToolResult { + t.Helper() + + ctx := context.Background() + result, err := client.CallTool(ctx, toolName, args) + require.NoError(t, err) + RequireNoError(t, result) + + return result +} + +// GetTextContent extracts text from the first content item. +func GetTextContent(t *testing.T, result *mcp.CallToolResult) string { + t.Helper() + require.NotEmpty(t, result.Content, "should have content in response") + + textContent, ok := result.Content[0].(*mcp.TextContent) + require.True(t, ok, "expected TextContent, got %T", result.Content[0]) + + return textContent.Text +}