tbphp
d4699dd2b1
* feat(health): 支持凭据按模型冷却与统一恢复 * fix(health): 修复模型冷却的错误契约与恢复边界 * fix(web): 修复模型冷却展示与重置反馈 * fix(web): 保持分组健康与恢复反馈一致 * fix(health): 修复停用分组的模型冷却身份展示
84 lines
4.0 KiB
Go
84 lines
4.0 KiB
Go
package gateway
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"gpt-load/internal/execution"
|
|
)
|
|
|
|
func TestModelCooldownHandlesSSERejectionAndCanceledResponse(t *testing.T) {
|
|
for _, canceled := range []bool{false, true} {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
result := UpstreamResult{StatusCode: 200, Header: http.Header{"Retry-After": {"3600"}},
|
|
Body: []byte(`{"error":{"type":"rate_limit_error"}}`), ProviderErrorBeforeCommit: true,
|
|
RequestWritten: true, DispatchState: execution.DispatchMaybeSent, ResponseStarted: true,
|
|
ExecutionError: &execution.ErrorEvidence{Kind: execution.ErrorKindProvider, Hint: execution.FailureHintRateLimited,
|
|
ScopeHint: execution.ErrorScopeModel, StatusCode: 429, ReplaySafety: execution.ReplaySafetyRejectedBeforeProcessing}}
|
|
forwarder := &scriptedForwarder{results: []UpstreamResult{result}}
|
|
if canceled {
|
|
forwarder.onCall = func(int) { cancel() }
|
|
}
|
|
engine, _, _, registry := newRequestLogHandlerTestRuntime(t, forwarder, &recordingAccessKeyRPMLimiter{}, &recordingRequestLogSink{}, "sk-one")
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o","messages":[]}`)).WithContext(ctx)
|
|
request.Header.Set("Authorization", "Bearer gl-client")
|
|
response := httptest.NewRecorder()
|
|
engine.ServeHTTP(response, request)
|
|
cancel()
|
|
if !canceled && response.Code != 429 {
|
|
t.Fatalf("SSE limit became %d", response.Code)
|
|
}
|
|
if len(registry.ModelCooldowns(1, time.Now())) != 1 {
|
|
t.Fatalf("canceled=%t: lost decoded quota rejection", canceled)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestModelCooldownRetriesAndFiltersFutureRequests(t *testing.T) {
|
|
forwarder := &scriptedForwarder{results: []UpstreamResult{
|
|
{StatusCode: 429, Header: http.Header{"Retry-After": {"3600"}}, Body: []byte(`{"error":{"code":"rate_limit_exceeded"}}`),
|
|
RequestWritten: true, DispatchState: execution.DispatchMaybeSent, ExecutionError: &execution.ErrorEvidence{
|
|
Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintRateLimited, StatusCode: 429}},
|
|
successfulAffinityResult(), successfulAffinityResult(),
|
|
}}
|
|
engine, _, _, registry := newRequestLogHandlerTestRuntime(t, forwarder, &recordingAccessKeyRPMLimiter{}, &recordingRequestLogSink{}, "sk-one", "sk-two")
|
|
for range 2 {
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o","messages":[]}`))
|
|
request.Header.Set("Authorization", "Bearer gl-client")
|
|
response := httptest.NewRecorder()
|
|
engine.ServeHTTP(response, request)
|
|
if response.Code != 200 {
|
|
t.Fatalf("response=%d %s", response.Code, response.Body.String())
|
|
}
|
|
}
|
|
if len(forwarder.inputs) != 3 || forwarder.inputs[0].Credential.ID == forwarder.inputs[1].Credential.ID || forwarder.inputs[0].Credential.ID == forwarder.inputs[2].Credential.ID {
|
|
t.Fatalf("unexpected candidate retry chain: %#v", forwarder.inputs)
|
|
}
|
|
limited := forwarder.inputs[0].Credential.ID
|
|
if len(registry.ModelCooldowns(limited, time.Now())) != 1 {
|
|
t.Fatal("no model cooldown")
|
|
}
|
|
if until, _ := registry.CredentialCooldownUntil(limited); !until.IsZero() {
|
|
t.Fatal("limited the whole credential")
|
|
}
|
|
}
|
|
|
|
func TestOnlyModelCooledCandidatesReturn429WithoutDispatch(t *testing.T) {
|
|
forwarder := &scriptedForwarder{}
|
|
engine, _, _, registry := newRequestLogHandlerTestRuntime(t, forwarder, &recordingAccessKeyRPMLimiter{}, &recordingRequestLogSink{}, "sk-one")
|
|
now := time.Now()
|
|
ref, _ := registry.CredentialRef(1)
|
|
registry.SetModelCooldown(ref, "gpt-4o", now.Add(time.Hour), now)
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o","messages":[]}`))
|
|
request.Header.Set("Authorization", "Bearer gl-client")
|
|
response := httptest.NewRecorder()
|
|
engine.ServeHTTP(response, request)
|
|
if response.Code != 429 || response.Header().Get("Retry-After") == "" || len(forwarder.inputs) != 0 {
|
|
t.Fatalf("limited candidate response = %d %s; attempts=%d", response.Code, response.Body.String(), len(forwarder.inputs))
|
|
}
|
|
}
|