test(mcp): auth middleware + rate limiter tests

2026-05-08 08:11:02 +08:00
parent e1e6508c4c
commit c046d77a6e
2 changed files with 204 additions and 0 deletions
@@ -0,0 +1,102 @@
+package mcp_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/yan1h/agent-coding-workflow/internal/mcp"
+)
+
+func TestAuthMiddleware_BearerHeader(t *testing.T) {
+	t.Parallel()
+	repo := newFakeRepo()
+	svc := mcp.NewTokenService(repo, &fakeAudit{})
+
+	res, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
+		UserID:    uuid.New(),
+		Name:      "x",
+		ExpiresAt: time.Now().Add(time.Hour),
+		CreatedBy: uuid.New(),
+	})
+	require.NoError(t, err)
+
+	mw := mcp.NewAuthMiddleware(svc)
+	h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		s, ok := mcp.FromContext(r.Context())
+		require.True(t, ok)
+		require.NotNil(t, s)
+		_, _ = w.Write([]byte(s.TokenID))
+	}))
+
+	req := httptest.NewRequest("GET", "/mcp", nil)
+	req.Header.Set("Authorization", "Bearer "+res.Plaintext)
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	body, _ := io.ReadAll(rec.Body)
+	assert.Equal(t, res.ID.String(), string(body))
+}
+
+func TestAuthMiddleware_QueryFallback(t *testing.T) {
+	t.Parallel()
+	repo := newFakeRepo()
+	svc := mcp.NewTokenService(repo, &fakeAudit{})
+
+	res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
+		UserID:    uuid.New(),
+		Name:      "x",
+		ExpiresAt: time.Now().Add(time.Hour),
+		CreatedBy: uuid.New(),
+	})
+
+	mw := mcp.NewAuthMiddleware(svc)
+	h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/mcp?token="+res.Plaintext, nil)
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestAuthMiddleware_MissingToken(t *testing.T) {
+	t.Parallel()
+	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
+
+	mw := mcp.NewAuthMiddleware(svc)
+	h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Fatal("should not call inner handler")
+	}))
+
+	req := httptest.NewRequest("GET", "/mcp", nil)
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusUnauthorized, rec.Code)
+}
+
+func TestAuthMiddleware_InvalidToken(t *testing.T) {
+	t.Parallel()
+	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
+
+	mw := mcp.NewAuthMiddleware(svc)
+	h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Fatal("should not call")
+	}))
+
+	req := httptest.NewRequest("GET", "/mcp", nil)
+	req.Header.Set("Authorization", "Bearer mcpt_invalid")
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusUnauthorized, rec.Code)
+}