Add pooled same-provider credential fallback
This commit is contained in:
@@ -1528,6 +1528,62 @@ class TestNousCredentialRefresh:
|
||||
assert isinstance(agent.client, _RebuiltClient)
|
||||
|
||||
|
||||
class TestCredentialPoolRecovery:
|
||||
def test_recover_with_pool_rotates_on_402(self, agent):
|
||||
current = SimpleNamespace(label="primary")
|
||||
next_entry = SimpleNamespace(label="secondary")
|
||||
|
||||
class _Pool:
|
||||
def current(self):
|
||||
return current
|
||||
|
||||
def mark_exhausted_and_rotate(self, *, status_code):
|
||||
assert status_code == 402
|
||||
return next_entry
|
||||
|
||||
agent._credential_pool = _Pool()
|
||||
agent._swap_credential = MagicMock()
|
||||
|
||||
recovered, retry_same = agent._recover_with_credential_pool(
|
||||
status_code=402,
|
||||
retry_429_with_same_cred=False,
|
||||
)
|
||||
|
||||
assert recovered is True
|
||||
assert retry_same is False
|
||||
agent._swap_credential.assert_called_once_with(next_entry)
|
||||
|
||||
def test_recover_with_pool_retries_first_429_then_rotates(self, agent):
|
||||
next_entry = SimpleNamespace(label="secondary")
|
||||
|
||||
class _Pool:
|
||||
def current(self):
|
||||
return SimpleNamespace(label="primary")
|
||||
|
||||
def mark_exhausted_and_rotate(self, *, status_code):
|
||||
assert status_code == 429
|
||||
return next_entry
|
||||
|
||||
agent._credential_pool = _Pool()
|
||||
agent._swap_credential = MagicMock()
|
||||
|
||||
recovered, retry_same = agent._recover_with_credential_pool(
|
||||
status_code=429,
|
||||
retry_429_with_same_cred=False,
|
||||
)
|
||||
assert recovered is False
|
||||
assert retry_same is True
|
||||
agent._swap_credential.assert_not_called()
|
||||
|
||||
recovered, retry_same = agent._recover_with_credential_pool(
|
||||
status_code=429,
|
||||
retry_429_with_same_cred=True,
|
||||
)
|
||||
assert recovered is True
|
||||
assert retry_same is False
|
||||
agent._swap_credential.assert_called_once_with(next_entry)
|
||||
|
||||
|
||||
class TestMaxTokensParam:
|
||||
"""Verify _max_tokens_param returns the correct key for each provider."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user