You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
38 lines
1.1 KiB
38 lines
1.1 KiB
import os
|
|
import time
|
|
|
|
import ai_provider
|
|
|
|
|
|
class DummyResponse:
|
|
def __init__(self, status_code=200, json_data=None, headers=None):
|
|
self.status_code = status_code
|
|
self._json = json_data or {}
|
|
self.headers = headers or {}
|
|
|
|
def json(self):
|
|
return self._json
|
|
|
|
|
|
def test_retry_on_429_then_success(monkeypatch):
|
|
calls = {"n": 0}
|
|
|
|
def fake_post(url, json, headers, timeout):
|
|
calls["n"] += 1
|
|
if calls["n"] <= 2:
|
|
# first two calls return 429 with Retry-After: 1
|
|
return DummyResponse(
|
|
429, json_data={"error": "rate_limited"}, headers={"Retry-After": "1"}
|
|
)
|
|
return DummyResponse(200, json_data={"data": [{"embedding": [0.4, 0.5]}]})
|
|
|
|
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-test")
|
|
monkeypatch.setattr("requests.post", fake_post)
|
|
|
|
start = time.time()
|
|
emb = ai_provider.get_embedding("hello")
|
|
duration = time.time() - start
|
|
|
|
# we should have waited at least ~2 seconds due to two Retry-After: 1 sleeps
|
|
assert duration >= 2
|
|
assert emb == [0.4, 0.5]
|
|
|