improve test coverage

Signed-off-by: Sanskarzz <sanskar.gur@gmail.com>

Author: Sanskarzz

pkg/ratelimit/middleware_test.go

@@ -13,11 +13,17 @@ import (
 	"testing"
 	"time"
 
+	"github.com/alicebob/miniredis/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/mcp"
+	transporttypes "github.com/stacklok/toolhive/pkg/transport/types"
+	transportmocks "github.com/stacklok/toolhive/pkg/transport/types/mocks"
 )
 
 // dummyLimiter is a test double for the Limiter interface.
@@ -208,3 +214,70 @@ func TestRateLimitHandler_NoIdentityPassesEmptyUserID(t *testing.T) {
 	assert.Equal(t, "echo", recorder.toolName)
 	assert.Empty(t, recorder.userID, "unauthenticated requests should pass empty userID")
 }
+
+func TestDefaultToolNameResolverNilParsedRequest(t *testing.T) {
+	t.Parallel()
+
+	assert.Empty(t, DefaultToolNameResolver(nil))
+}
+
+func TestNewMiddlewareUsesCustomToolNameResolver(t *testing.T) {
+	t.Parallel()
+
+	recorder := &recordingLimiter{}
+	handler := NewMiddleware(recorder, func(*mcp.ParsedMCPRequest) string {
+		return "resolved-tool"
+	})(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
+	req = withParsedMCPRequest(req, "tools/call", "raw-tool", 1)
+	w := httptest.NewRecorder()
+
+	handler.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	assert.Equal(t, "resolved-tool", recorder.toolName)
+}
+
+func TestRateLimitMiddlewareHandlerReturnsConfiguredHandler(t *testing.T) {
+	t.Parallel()
+
+	expected := rateLimitHandler(&dummyLimiter{decision: &Decision{Allowed: true}})
+	mw := &rateLimitMiddleware{handler: expected}
+
+	assert.NotNil(t, mw.Handler())
+}
+
+func TestCreateMiddlewareRegistersUsableMiddleware(t *testing.T) {
+	t.Parallel()
+
+	mr := miniredis.RunT(t)
+	cfg, err := transporttypes.NewMiddlewareConfig(MiddlewareType, MiddlewareParams{
+		Namespace:  "default",
+		ServerName: "server",
+		RedisAddr:  mr.Addr(),
+		Config: &v1beta1.RateLimitConfig{
+			Shared: &v1beta1.RateLimitBucket{
+				MaxTokens:    1,
+				RefillPeriod: metav1.Duration{Duration: time.Minute},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	ctrl := gomock.NewController(t)
+	runner := transportmocks.NewMockMiddlewareRunner(ctrl)
+	var registered transporttypes.Middleware
+	runner.EXPECT().
+		AddMiddleware(MiddlewareType, gomock.AssignableToTypeOf(&rateLimitMiddleware{})).
+		Do(func(_ string, middleware transporttypes.Middleware) {
+			registered = middleware
+		})
+
+	require.NoError(t, CreateMiddleware(cfg, runner))
+	require.NotNil(t, registered)
+	require.NotNil(t, registered.Handler())
+	require.NoError(t, registered.Close())
+}

pkg/vmcp/cli/serve_test.go

@@ -337,6 +337,20 @@ func TestValidateQuickModeHost(t *testing.T) {
 	}
 }
 
+func TestVMCPNamespace(t *testing.T) {
+	t.Run("defaults to local", func(t *testing.T) {
+		t.Setenv("VMCP_NAMESPACE", "")
+
+		assert.Equal(t, "local", vmcpNamespace())
+	})
+
+	t.Run("uses environment value", func(t *testing.T) {
+		t.Setenv("VMCP_NAMESPACE", "toolhive-system")
+
+		assert.Equal(t, "toolhive-system", vmcpNamespace())
+	})
+}
+
 // TestRunDiscovery_ZeroBackends exercises the branch in runDiscovery where the
 // discoverer succeeds but returns no backends. The function must return a
 // non-error, an empty (non-nil) backend slice, and pass through the client and

pkg/vmcp/server/ratelimit_test.go

@@ -15,13 +15,18 @@ import (
 	"github.com/alicebob/miniredis/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/stacklok/toolhive/pkg/auth"
 	mcpparser "github.com/stacklok/toolhive/pkg/mcp"
 	ratelimittypes "github.com/stacklok/toolhive/pkg/ratelimit/types"
+	"github.com/stacklok/toolhive/pkg/vmcp"
 	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
+	discoveryMocks "github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks"
+	"github.com/stacklok/toolhive/pkg/vmcp/mocks"
 	"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
+	routerMocks "github.com/stacklok/toolhive/pkg/vmcp/router/mocks"
 )
 
 func TestBuildRateLimitMiddlewareDisabledWithoutConfig(t *testing.T) {
@@ -55,6 +60,82 @@ func TestBuildRateLimitMiddlewareRequiresRedisSessionStorage(t *testing.T) {
 	assert.Nil(t, cleanup)
 }
 
+func TestBuildRateLimitMiddlewareRequiresRedisAddress(t *testing.T) {
+	t.Parallel()
+
+	s := &Server{
+		config: &Config{
+			Name:         "vmcp",
+			Namespace:    "default",
+			RateLimiting: sharedRateLimitConfig(1),
+			SessionStorage: &vmcpconfig.SessionStorageConfig{
+				Provider: "redis",
+			},
+		},
+	}
+
+	middleware, cleanup, err := s.buildRateLimitMiddleware(t.Context())
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "requires Redis session storage address")
+	assert.Nil(t, middleware)
+	assert.Nil(t, cleanup)
+}
+
+func TestBuildRateLimitMiddlewareRedisPingFailure(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
+	defer cancel()
+	s := &Server{
+		config: &Config{
+			Name:         "vmcp",
+			Namespace:    "default",
+			RateLimiting: sharedRateLimitConfig(1),
+			SessionStorage: &vmcpconfig.SessionStorageConfig{
+				Provider: "redis",
+				Address:  "127.0.0.1:1",
+			},
+		},
+	}
+
+	middleware, cleanup, err := s.buildRateLimitMiddleware(ctx)
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to connect to Redis")
+	assert.Nil(t, middleware)
+	assert.Nil(t, cleanup)
+}
+
+func TestBuildRateLimitMiddlewareInvalidRateLimitConfig(t *testing.T) {
+	t.Parallel()
+
+	mr := miniredis.RunT(t)
+	s := &Server{
+		config: &Config{
+			Name:      "vmcp",
+			Namespace: "default",
+			RateLimiting: &ratelimittypes.RateLimitConfig{
+				Shared: &ratelimittypes.RateLimitBucket{
+					MaxTokens:    0,
+					RefillPeriod: metav1.Duration{Duration: time.Minute},
+				},
+			},
+			SessionStorage: &vmcpconfig.SessionStorageConfig{
+				Provider: "redis",
+				Address:  mr.Addr(),
+			},
+		},
+	}
+
+	middleware, cleanup, err := s.buildRateLimitMiddleware(t.Context())
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to create rate limiter")
+	assert.Nil(t, middleware)
+	assert.Nil(t, cleanup)
+}
+
 func TestRateLimitMiddlewarePerUserSharedAcrossTools(t *testing.T) {
 	t.Parallel()
 
@@ -150,6 +231,96 @@ func TestRateLimitToolNameFallsBackToCallTool(t *testing.T) {
 	assert.Equal(t, "call_tool", s.rateLimitToolName(parsed))
 }
 
+func TestRateLimitToolNameNilParsedRequest(t *testing.T) {
+	t.Parallel()
+
+	s := &Server{config: &Config{}}
+
+	assert.Empty(t, s.rateLimitToolName(nil))
+}
+
+func TestRateLimitToolNameOptimizerFallsBackForInvalidInnerToolName(t *testing.T) {
+	t.Parallel()
+
+	s := &Server{config: &Config{OptimizerConfig: &optimizer.Config{}}}
+	parsed := &mcpparser.ParsedMCPRequest{
+		Method:     "tools/call",
+		ResourceID: "call_tool",
+		Arguments:  map[string]any{"tool_name": 123},
+	}
+
+	assert.Equal(t, "call_tool", s.rateLimitToolName(parsed))
+}
+
+func TestApplyRateLimitingWrapsConfiguredMiddleware(t *testing.T) {
+	t.Parallel()
+
+	s := &Server{
+		rateLimitMiddleware: func(next http.Handler) http.Handler {
+			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("X-Rate-Limit-Test", "wrapped")
+				next.ServeHTTP(w, r)
+			})
+		},
+	}
+	handler := s.applyRateLimiting(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusAccepted)
+	}))
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusAccepted, rec.Code)
+	assert.Equal(t, "wrapped", rec.Header().Get("X-Rate-Limit-Test"))
+}
+
+func TestApplyAuthorizationWrapsConfiguredMiddleware(t *testing.T) {
+	t.Parallel()
+
+	s := &Server{config: &Config{
+		AuthzMiddleware: func(next http.Handler) http.Handler {
+			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("X-Authz-Test", "wrapped")
+				next.ServeHTTP(w, r)
+			})
+		},
+	}}
+	handler := s.applyAuthorization(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusAccepted)
+	}))
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusAccepted, rec.Code)
+	assert.Equal(t, "wrapped", rec.Header().Get("X-Authz-Test"))
+}
+
+func TestNewDefaultsNamespace(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+	mockRouter := routerMocks.NewMockRouter(ctrl)
+	mockBackendClient := mocks.NewMockBackendClient(ctrl)
+	mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
+
+	s, err := New(
+		t.Context(),
+		&Config{SessionFactory: testMinimalFactory()},
+		mockRouter,
+		mockBackendClient,
+		mockDiscoveryMgr,
+		vmcp.NewImmutableRegistry([]vmcp.Backend{}),
+		nil,
+	)
+
+	require.NoError(t, err)
+	assert.Equal(t, "local", s.config.Namespace)
+}
+
 func newTestRateLimitHandler(t *testing.T, cfg *Config) http.Handler {
 	t.Helper()