| 392 | } |
| 393 | |
| 394 | func TestMarkKeyOnError(t *testing.T) { |
| 395 | t.Parallel() |
| 396 | |
| 397 | tests := []struct { |
| 398 | name string |
| 399 | err error |
| 400 | expectedReturn bool |
| 401 | expectedState keypool.KeyState |
| 402 | }{ |
| 403 | { |
| 404 | // Not an *openai.Error: no status code to act on. |
| 405 | name: "non_api_error_returns_false", |
| 406 | err: xerrors.New("network failure"), |
| 407 | expectedReturn: false, |
| 408 | expectedState: keypool.KeyStateValid, |
| 409 | }, |
| 410 | { |
| 411 | // Rate-limited: temporary cooldown. |
| 412 | name: "429_marks_temporary", |
| 413 | err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}}, |
| 414 | expectedReturn: true, |
| 415 | expectedState: keypool.KeyStateTemporary, |
| 416 | }, |
| 417 | { |
| 418 | // Auth failure: mark permanent. |
| 419 | name: "401_marks_permanent", |
| 420 | err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, |
| 421 | expectedReturn: true, |
| 422 | expectedState: keypool.KeyStatePermanent, |
| 423 | }, |
| 424 | { |
| 425 | // Auth forbidden: mark permanent. |
| 426 | name: "403_marks_permanent", |
| 427 | err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, |
| 428 | expectedReturn: true, |
| 429 | expectedState: keypool.KeyStatePermanent, |
| 430 | }, |
| 431 | { |
| 432 | // Server errors are not key-specific. |
| 433 | name: "500_does_not_mark", |
| 434 | err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}}, |
| 435 | expectedReturn: false, |
| 436 | expectedState: keypool.KeyStateValid, |
| 437 | }, |
| 438 | } |
| 439 | |
| 440 | for _, tc := range tests { |
| 441 | t.Run(tc.name, func(t *testing.T) { |
| 442 | t.Parallel() |
| 443 | pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t)) |
| 444 | require.NoError(t, err) |
| 445 | key, keyPoolErr := pool.Walker().Next() |
| 446 | require.Nil(t, keyPoolErr) |
| 447 | |
| 448 | base := &responsesInterceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()} |
| 449 | |
| 450 | got := base.markKeyOnError(context.Background(), key, tc.err) |
| 451 | assert.Equal(t, tc.expectedReturn, got) |