128 lines
3.3 KiB
Go
128 lines
3.3 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type modelNotFoundRateLimitCall struct {
|
|
accountID int64
|
|
scope string
|
|
resetAt time.Time
|
|
reason string
|
|
}
|
|
|
|
type modelNotFoundAccountRepoStub struct {
|
|
mockAccountRepoForGemini
|
|
tempCalls int
|
|
modelRateLimitCalls []modelNotFoundRateLimitCall
|
|
modelRateLimitErr error
|
|
}
|
|
|
|
func (r *modelNotFoundAccountRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
|
r.tempCalls++
|
|
return nil
|
|
}
|
|
|
|
func (r *modelNotFoundAccountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
|
|
call := modelNotFoundRateLimitCall{
|
|
accountID: id,
|
|
scope: scope,
|
|
resetAt: resetAt,
|
|
}
|
|
if len(reason) > 0 {
|
|
call.reason = reason[0]
|
|
}
|
|
r.modelRateLimitCalls = append(r.modelRateLimitCalls, call)
|
|
return r.modelRateLimitErr
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_ModelNotFoundUsesModelRateLimit(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := openAIModelNotFoundTempAccount()
|
|
|
|
handled := svc.HandleUpstreamError(
|
|
context.Background(),
|
|
account,
|
|
http.StatusNotFound,
|
|
http.Header{},
|
|
[]byte(`{"error":{"code":"model_not_found","message":"model not found"}}`),
|
|
"gpt-5.4",
|
|
)
|
|
|
|
require.True(t, handled)
|
|
require.Zero(t, repo.tempCalls)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
call := repo.modelRateLimitCalls[0]
|
|
require.Equal(t, account.ID, call.accountID)
|
|
require.Equal(t, "gpt-5.4", call.scope)
|
|
require.Equal(t, upstreamModelNotFoundReason, call.reason)
|
|
require.WithinDuration(t, time.Now().Add(upstreamModelNotFoundCooldown), call.resetAt, 5*time.Second)
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_ModelNotFoundWriteFailureDoesNotTempUnschedule(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := openAIModelNotFoundTempAccount()
|
|
|
|
handled := svc.HandleUpstreamError(
|
|
context.Background(),
|
|
account,
|
|
http.StatusNotFound,
|
|
http.Header{},
|
|
[]byte(`{"error":{"code":"model_not_found","message":"model not found"}}`),
|
|
"gpt-5.4",
|
|
)
|
|
|
|
require.True(t, handled)
|
|
require.Zero(t, repo.tempCalls)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_Bare404KeepsTempUnschedulablePath(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := openAIModelNotFoundTempAccount()
|
|
|
|
handled := svc.HandleUpstreamError(
|
|
context.Background(),
|
|
account,
|
|
http.StatusNotFound,
|
|
http.Header{},
|
|
[]byte(`{"error":{"message":"endpoint not found"}}`),
|
|
"gpt-5.4",
|
|
)
|
|
|
|
require.True(t, handled)
|
|
require.Equal(t, 1, repo.tempCalls)
|
|
require.Empty(t, repo.modelRateLimitCalls)
|
|
}
|
|
|
|
func openAIModelNotFoundTempAccount() *Account {
|
|
return &Account{
|
|
ID: 101,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
Status: StatusActive,
|
|
Schedulable: true,
|
|
Credentials: map[string]any{
|
|
"temp_unschedulable_enabled": true,
|
|
"temp_unschedulable_rules": []any{
|
|
map[string]any{
|
|
"error_code": float64(http.StatusNotFound),
|
|
"keywords": []any{"not found"},
|
|
"duration_minutes": float64(10),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|