65 lines
1.8 KiB
Go
65 lines
1.8 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestRateLimitService_HandleUpstreamError_OpenAI403FirstHitTempUnschedulable(t *testing.T) {
|
|
repo := &rateLimitAccountRepoStub{}
|
|
counter := &openAI403CounterCacheStub{counts: []int64{1}}
|
|
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
|
service.SetOpenAI403CounterCache(counter)
|
|
account := &Account{
|
|
ID: 301,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
}
|
|
|
|
shouldDisable := service.HandleUpstreamError(
|
|
context.Background(),
|
|
account,
|
|
http.StatusForbidden,
|
|
http.Header{},
|
|
[]byte(`{"error":{"message":"temporary edge rejection"}}`),
|
|
)
|
|
|
|
require.True(t, shouldDisable)
|
|
require.Equal(t, 0, repo.setErrorCalls)
|
|
require.Equal(t, 1, repo.tempCalls)
|
|
require.Contains(t, repo.lastTempReason, "temporary edge rejection")
|
|
require.Contains(t, repo.lastTempReason, "(1/3)")
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_OpenAI403ThresholdDisables(t *testing.T) {
|
|
repo := &rateLimitAccountRepoStub{}
|
|
counter := &openAI403CounterCacheStub{counts: []int64{3}}
|
|
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
|
service.SetOpenAI403CounterCache(counter)
|
|
account := &Account{
|
|
ID: 302,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
}
|
|
|
|
shouldDisable := service.HandleUpstreamError(
|
|
context.Background(),
|
|
account,
|
|
http.StatusForbidden,
|
|
http.Header{},
|
|
[]byte(`{"error":{"message":"workspace forbidden by policy"}}`),
|
|
)
|
|
|
|
require.True(t, shouldDisable)
|
|
require.Equal(t, 1, repo.setErrorCalls)
|
|
require.Equal(t, 0, repo.tempCalls)
|
|
require.Contains(t, repo.lastErrorMsg, "workspace forbidden by policy")
|
|
require.Contains(t, repo.lastErrorMsg, "consecutive_403=3/3")
|
|
}
|