You've already forked agentic-coding-workflow
103 lines
2.7 KiB
Go
103 lines
2.7 KiB
Go
package mcp_test
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/mcp"
|
|
)
|
|
|
|
func TestRateLimiter_AllowUntilCap(t *testing.T) {
|
|
t.Parallel()
|
|
now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
|
|
limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
|
|
PerTokenPerMinute: 3,
|
|
GlobalPerMinute: 100,
|
|
}, func() time.Time { return now })
|
|
|
|
for i := 0; i < 3; i++ {
|
|
ok, _ := limiter.Allow("token-A")
|
|
assert.True(t, ok, "request %d should pass", i)
|
|
}
|
|
ok, reason := limiter.Allow("token-A")
|
|
assert.False(t, ok)
|
|
assert.Equal(t, "per_token", reason)
|
|
}
|
|
|
|
func TestRateLimiter_GlobalCap(t *testing.T) {
|
|
t.Parallel()
|
|
now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
|
|
limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
|
|
PerTokenPerMinute: 100,
|
|
GlobalPerMinute: 2,
|
|
}, func() time.Time { return now })
|
|
|
|
ok, _ := limiter.Allow("token-A")
|
|
assert.True(t, ok)
|
|
ok, _ = limiter.Allow("token-B")
|
|
assert.True(t, ok)
|
|
ok, reason := limiter.Allow("token-C")
|
|
assert.False(t, ok)
|
|
assert.Equal(t, "global", reason)
|
|
}
|
|
|
|
func TestRateLimiter_WindowReset(t *testing.T) {
|
|
t.Parallel()
|
|
current := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
|
|
limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
|
|
PerTokenPerMinute: 1,
|
|
GlobalPerMinute: 100,
|
|
}, func() time.Time { return current })
|
|
|
|
ok, _ := limiter.Allow("A")
|
|
assert.True(t, ok)
|
|
ok, _ = limiter.Allow("A")
|
|
assert.False(t, ok)
|
|
|
|
// 跨分钟窗口
|
|
current = current.Add(time.Minute)
|
|
ok, _ = limiter.Allow("A")
|
|
assert.True(t, ok, "after window reset should allow again")
|
|
}
|
|
|
|
func TestRateLimitMiddleware_429(t *testing.T) {
|
|
t.Parallel()
|
|
now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
|
|
limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
|
|
PerTokenPerMinute: 1,
|
|
GlobalPerMinute: 100,
|
|
}, func() time.Time { return now })
|
|
|
|
// 内层 handler 计数被调次数
|
|
var hits atomic.Int32
|
|
innerH := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
hits.Add(1)
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
// 模拟 auth 中间件已注入 ctx：直接 context.WithValue + mcp.AuthContextKey
|
|
withSession := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ctx := context.WithValue(r.Context(), mcp.AuthContextKey, &mcp.AuthSession{
|
|
TokenID: "fake-token-id",
|
|
UserID: "fake-user",
|
|
})
|
|
mcp.NewRateLimitMiddleware(limiter)(innerH).ServeHTTP(w, r.WithContext(ctx))
|
|
})
|
|
|
|
req := httptest.NewRequest("GET", "/mcp", nil)
|
|
rec := httptest.NewRecorder()
|
|
withSession.ServeHTTP(rec, req)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
rec2 := httptest.NewRecorder()
|
|
withSession.ServeHTTP(rec2, req)
|
|
assert.Equal(t, http.StatusTooManyRequests, rec2.Code)
|
|
assert.Equal(t, int32(1), hits.Load())
|
|
}
|