Function TestResponseErrorFromKeyPool

aibridge/intercept/openai_errors_test.go:15–55

(t *testing.T)

)

func TestResponseErrorFromKeyPool(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name               string
		keyPoolErr         *keypool.Error
		expectedStatus     int
		expectedRetryAfter time.Duration
	}{
		{
			// Rate-limited with no cooldown: 429, no Retry-After.
			name:               "rate_limited_zero_retry_after",
			keyPoolErr:         &keypool.Error{Kind: keypool.ErrorKindRateLimited},
			expectedStatus:     http.StatusTooManyRequests,
			expectedRetryAfter: 0,
		},
		{
			// Rate-limited with cooldown: 429, Retry-After set.
			name:               "rate_limited_with_retry_after",
			keyPoolErr:         &keypool.Error{Kind: keypool.ErrorKindRateLimited, RetryAfter: 5 * time.Second},
			expectedStatus:     http.StatusTooManyRequests,
			expectedRetryAfter: 5 * time.Second,
		},
		{
			// Permanent: 502 api_error.
			name:           "permanent_returns_502",
			keyPoolErr:     &keypool.Error{Kind: keypool.ErrorKindPermanent},
			expectedStatus: http.StatusBadGateway,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got := intercept.ResponseErrorFromKeyPool(tc.keyPoolErr)
			require.NotNil(t, got)
			assert.Equal(t, tc.expectedStatus, got.StatusCode)
			assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter)
		})
	}
}