superidou
5759c1862e
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
84 lines
3.4 KiB
Go
84 lines
3.4 KiB
Go
package shadow
|
|
|
|
import (
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"aigateway.local/core/internal/gateway"
|
|
"aigateway.local/core/internal/platform/config"
|
|
)
|
|
|
|
func TestContractSignatureComparesShapeNotGeneratedValues(t *testing.T) {
|
|
left := contractSignature(200, []byte(`{"id":"one","choices":[{"message":{"content":"answer one"}}]}`))
|
|
right := contractSignature(200, []byte(`{"id":"two","choices":[{"message":{"content":"answer two"}}]}`))
|
|
if left != right {
|
|
t.Fatalf("expected equal shapes: %s != %s", left, right)
|
|
}
|
|
if left == contractSignature(500, []byte(`{"choices":[]}`)) {
|
|
t.Fatal("status and schema mismatch must differ")
|
|
}
|
|
}
|
|
|
|
func TestMiddlewareUsesDedicatedCredentialAndDoesNotDelayPrimary(t *testing.T) {
|
|
delivered := make(chan string, 1)
|
|
legacy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
raw, _ := io.ReadAll(r.Body)
|
|
delivered <- r.Header.Get("Authorization") + ":" + string(raw)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"id":"legacy","choices":[{"message":{"content":"different"}}]}`))
|
|
}))
|
|
defer legacy.Close()
|
|
middleware := New(config.Shadow{BaseURL: legacy.URL, APIKey: "shadow-secret", SampleRate: 1, Timeout: time.Second, MaxBodyBytes: 1 << 20, MaxConcurrent: 2}, slog.Default())
|
|
primary := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"id":"go","choices":[{"message":{"content":"primary"}}]}`))
|
|
})
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","stream":false}`))
|
|
request.Header.Set("Authorization", "Bearer production-secret")
|
|
request = request.WithContext(gateway.WithRequestID(request.Context(), "shadow-test"))
|
|
response := httptest.NewRecorder()
|
|
middleware.Wrap(primary).ServeHTTP(response, request)
|
|
if response.Code != 200 {
|
|
t.Fatalf("primary status %d", response.Code)
|
|
}
|
|
select {
|
|
case value := <-delivered:
|
|
if !strings.HasPrefix(value, "Bearer shadow-secret:") {
|
|
t.Fatalf("shadow credential not replaced: %s", value)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("shadow request was not delivered")
|
|
}
|
|
deadline := time.Now().Add(time.Second)
|
|
for middleware.metrics.matched.Load()+middleware.metrics.mismatched.Load() == 0 && time.Now().Before(deadline) {
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
if middleware.metrics.matched.Load() != 1 {
|
|
t.Fatalf("expected shape match, metrics=%s", middleware.Prometheus())
|
|
}
|
|
}
|
|
|
|
func TestOversizedShadowBodyIsPreservedForPrimary(t *testing.T) {
|
|
middleware := New(config.Shadow{BaseURL: "http://127.0.0.1:1", APIKey: "shadow", SampleRate: 1, Timeout: time.Second, MaxBodyBytes: 16, MaxConcurrent: 1}, slog.Default())
|
|
payload := `{"stream":false,"prompt":"this body is larger than the shadow limit"}`
|
|
primary := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
raw, _ := io.ReadAll(r.Body)
|
|
if string(raw) != payload {
|
|
t.Fatalf("primary body changed: %q", raw)
|
|
}
|
|
w.WriteHeader(http.StatusNoContent)
|
|
})
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(payload))
|
|
request = request.WithContext(gateway.WithRequestID(request.Context(), "large-shadow"))
|
|
response := httptest.NewRecorder()
|
|
middleware.Wrap(primary).ServeHTTP(response, request)
|
|
if response.Code != http.StatusNoContent {
|
|
t.Fatalf("status %d", response.Code)
|
|
}
|
|
}
|