| 88 | } |
| 89 | |
| 90 | func TestMarkKeyOnError(t *testing.T) { |
| 91 | t.Parallel() |
| 92 | |
| 93 | tests := []struct { |
| 94 | name string |
| 95 | err error |
| 96 | expectedReturn bool |
| 97 | expectedState keypool.KeyState |
| 98 | }{ |
| 99 | { |
| 100 | // Not an *openai.Error: no status code to act on. |
| 101 | name: "non_api_error_returns_false", |
| 102 | err: xerrors.New("network failure"), |
| 103 | expectedReturn: false, |
| 104 | expectedState: keypool.KeyStateValid, |
| 105 | }, |
| 106 | { |
| 107 | // Rate-limited: temporary cooldown. |
| 108 | name: "429_marks_temporary", |
| 109 | err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}}, |
| 110 | expectedReturn: true, |
| 111 | expectedState: keypool.KeyStateTemporary, |
| 112 | }, |
| 113 | { |
| 114 | // Auth failure: mark permanent. |
| 115 | name: "401_marks_permanent", |
| 116 | err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, |
| 117 | expectedReturn: true, |
| 118 | expectedState: keypool.KeyStatePermanent, |
| 119 | }, |
| 120 | { |
| 121 | // Auth forbidden: mark permanent. |
| 122 | name: "403_marks_permanent", |
| 123 | err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, |
| 124 | expectedReturn: true, |
| 125 | expectedState: keypool.KeyStatePermanent, |
| 126 | }, |
| 127 | { |
| 128 | // Server errors are not key-specific. |
| 129 | name: "500_does_not_mark", |
| 130 | err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}}, |
| 131 | expectedReturn: false, |
| 132 | expectedState: keypool.KeyStateValid, |
| 133 | }, |
| 134 | } |
| 135 | |
| 136 | for _, tc := range tests { |
| 137 | t.Run(tc.name, func(t *testing.T) { |
| 138 | t.Parallel() |
| 139 | pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t)) |
| 140 | require.NoError(t, err) |
| 141 | key, keyPoolErr := pool.Walker().Next() |
| 142 | require.Nil(t, keyPoolErr) |
| 143 | |
| 144 | base := &interceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()} |
| 145 | |
| 146 | got := base.markKeyOnError(context.Background(), key, tc.err) |
| 147 | assert.Equal(t, tc.expectedReturn, got) |