235 lines
8.8 KiB
Python
235 lines
8.8 KiB
Python
import pytest
|
|
from core import extractor
|
|
from core.extractor import (
|
|
extract_with_openrouter, ExtractorError, RateLimitError, CreditExhaustedError, _parse_json_response,
|
|
)
|
|
|
|
|
|
def test_parse_json_response_handles_clean_json():
|
|
result = _parse_json_response('{"tripName": "Peru Classic"}', 'model-a', 'test')
|
|
assert result == {'tripName': 'Peru Classic'}
|
|
|
|
|
|
def test_parse_json_response_strips_code_fences():
|
|
result = _parse_json_response('```json\n{"tripName": "Peru Classic"}\n```', 'model-a', 'test')
|
|
assert result == {'tripName': 'Peru Classic'}
|
|
|
|
|
|
def test_parse_json_response_extracts_object_from_surrounding_text():
|
|
result = _parse_json_response('Here is the data: {"tripName": "Peru Classic"} Thanks!', 'model-a', 'test')
|
|
assert result == {'tripName': 'Peru Classic'}
|
|
|
|
|
|
def test_parse_json_response_repairs_trailing_commas_and_unquoted_keys():
|
|
malformed = '{tripName: "Peru Classic", duration: 15,}'
|
|
result = _parse_json_response(malformed, 'model-a', 'test')
|
|
assert result == {'tripName': 'Peru Classic', 'duration': 15}
|
|
|
|
|
|
def test_parse_json_response_raises_when_unparseable():
|
|
with pytest.raises(ValueError):
|
|
_parse_json_response('not json at all !!!', 'model-a', 'test')
|
|
|
|
|
|
class _FakeChoice:
|
|
def __init__(self, content):
|
|
self.message = type('Msg', (), {'content': content})()
|
|
|
|
|
|
class _FakeCompletionResponse:
|
|
def __init__(self, content):
|
|
self.choices = [_FakeChoice(content)]
|
|
|
|
|
|
def test_call_openrouter_raises_rate_limit_error(monkeypatch):
|
|
class _FakeClient:
|
|
class chat:
|
|
class completions:
|
|
@staticmethod
|
|
def create(**kwargs):
|
|
raise Exception('429 Too Many Requests: rate_limit_exceeded')
|
|
|
|
monkeypatch.setattr(extractor, 'OpenAI', lambda api_key, base_url: _FakeClient())
|
|
with pytest.raises(RateLimitError):
|
|
extractor._call_openrouter('some-model', 'prompt text', '')
|
|
|
|
|
|
def test_call_openrouter_raises_credit_exhausted_error(monkeypatch):
|
|
class _FakeClient:
|
|
class chat:
|
|
class completions:
|
|
@staticmethod
|
|
def create(**kwargs):
|
|
raise Exception('403 Forbidden: insufficient credit balance for this model')
|
|
|
|
monkeypatch.setattr(extractor, 'OpenAI', lambda api_key, base_url: _FakeClient())
|
|
with pytest.raises(CreditExhaustedError):
|
|
extractor._call_openrouter('some-model', 'prompt text', '')
|
|
|
|
|
|
def test_call_openrouter_reraises_other_errors(monkeypatch):
|
|
class _FakeClient:
|
|
class chat:
|
|
class completions:
|
|
@staticmethod
|
|
def create(**kwargs):
|
|
raise Exception('500 server error')
|
|
|
|
monkeypatch.setattr(extractor, 'OpenAI', lambda api_key, base_url: _FakeClient())
|
|
with pytest.raises(Exception, match='500 server error'):
|
|
extractor._call_openrouter('some-model', 'prompt text', '')
|
|
|
|
|
|
def test_call_openrouter_appends_content_when_present(monkeypatch):
|
|
captured = {}
|
|
|
|
class _FakeClient:
|
|
class chat:
|
|
class completions:
|
|
@staticmethod
|
|
def create(**kwargs):
|
|
captured['messages'] = kwargs['messages']
|
|
return _FakeCompletionResponse('{"ok": true}')
|
|
|
|
monkeypatch.setattr(extractor, 'OpenAI', lambda api_key, base_url: _FakeClient())
|
|
result = extractor._call_openrouter('some-model', 'BASE PROMPT', 'EXTRA CONTENT')
|
|
|
|
assert result == {'ok': True}
|
|
assert 'BASE PROMPT' in captured['messages'][0]['content']
|
|
assert 'EXTRA CONTENT' in captured['messages'][0]['content']
|
|
|
|
|
|
def test_call_openrouter_forces_json_object_response_format(monkeypatch):
|
|
"""
|
|
Regression test — deepseek/deepseek-v4-flash wrote a full prose
|
|
comparative analysis instead of the requested JSON object on a real
|
|
production call, whenever the competitor's content didn't clearly
|
|
match the requested trip. response_format=json_object constrains
|
|
this at the API level rather than relying on prompt wording alone.
|
|
"""
|
|
captured = {}
|
|
|
|
class _FakeClient:
|
|
class chat:
|
|
class completions:
|
|
@staticmethod
|
|
def create(**kwargs):
|
|
captured.update(kwargs)
|
|
return _FakeCompletionResponse('{"ok": true}')
|
|
|
|
monkeypatch.setattr(extractor, 'OpenAI', lambda api_key, base_url: _FakeClient())
|
|
extractor._call_openrouter('some-model', 'prompt', '')
|
|
|
|
assert captured['response_format'] == {'type': 'json_object'}
|
|
|
|
|
|
def test_extract_with_openrouter_retries_same_model_on_rate_limit_before_moving_on(monkeypatch):
|
|
"""
|
|
Regression test — the Phase 2 refactor moved to the next model on
|
|
the very first rate-limit response, dropping the legacy app.py's
|
|
retry-with-backoff behavior. A model-free-1 that's merely rate
|
|
limited (not exhausted) should get RETRY_ATTEMPTS_PER_MODEL tries
|
|
before extract_with_openrouter gives up on it.
|
|
"""
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1', 'model-free-2'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
sleeps = []
|
|
monkeypatch.setattr(extractor.time, 'sleep', lambda seconds: sleeps.append(seconds))
|
|
|
|
calls = []
|
|
|
|
def _fake_call(model, prompt, content):
|
|
calls.append(model)
|
|
if model == 'model-free-1':
|
|
raise RateLimitError('rate limited')
|
|
return {'tripName': 'ok'}
|
|
|
|
monkeypatch.setattr(extractor, '_call_openrouter', _fake_call)
|
|
|
|
result = extract_with_openrouter('prompt')
|
|
|
|
assert result == {'tripName': 'ok'}
|
|
# Retried model-free-1 up to the attempt limit before moving to model-free-2.
|
|
assert calls == ['model-free-1', 'model-free-1', 'model-free-2']
|
|
assert sleeps == [5, 10] # 5 * (attempt + 1) for attempts 0 and 1
|
|
|
|
|
|
def test_extract_with_openrouter_moves_on_immediately_on_credit_exhausted(monkeypatch):
|
|
"""CreditExhaustedError shouldn't retry the same model — the quota won't recover mid-request."""
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1', 'model-free-2'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
|
|
calls = []
|
|
|
|
def _fake_call(model, prompt, content):
|
|
calls.append(model)
|
|
if model == 'model-free-1':
|
|
raise CreditExhaustedError('quota exhausted')
|
|
return {'tripName': 'ok'}
|
|
|
|
monkeypatch.setattr(extractor, '_call_openrouter', _fake_call)
|
|
|
|
result = extract_with_openrouter('prompt')
|
|
|
|
assert result == {'tripName': 'ok'}
|
|
assert calls == ['model-free-1', 'model-free-2'] # only one attempt on model-free-1
|
|
|
|
|
|
def test_extract_with_openrouter_retries_same_model_on_generic_error_before_moving_on(monkeypatch):
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
|
|
calls = []
|
|
|
|
def _fake_call(model, prompt, content):
|
|
calls.append(model)
|
|
if model == 'model-free-1':
|
|
raise Exception('exhausted')
|
|
return {'tripName': 'from-fallback'}
|
|
|
|
monkeypatch.setattr(extractor, '_call_openrouter', _fake_call)
|
|
sentry_calls = []
|
|
monkeypatch.setattr(extractor.sentry_sdk, 'capture_message', lambda msg, level=None: sentry_calls.append(msg))
|
|
|
|
result = extract_with_openrouter('prompt')
|
|
|
|
assert result == {'tripName': 'from-fallback'}
|
|
assert calls == ['model-free-1', 'model-free-1', 'model-paid']
|
|
assert sentry_calls == ['OpenRouter paid fallback used']
|
|
|
|
|
|
def test_extract_with_openrouter_raises_when_all_models_exhausted(monkeypatch):
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
monkeypatch.setattr(extractor, '_call_openrouter', lambda *a, **k: (_ for _ in ()).throw(Exception('down')))
|
|
|
|
with pytest.raises(ExtractorError):
|
|
extract_with_openrouter('prompt')
|
|
|
|
|
|
def test_extract_with_openrouter_skips_free_models_when_disabled(monkeypatch):
|
|
"""
|
|
MVP launch toggle — OPENROUTER_USE_FREE_MODELS_FIRST=false must call
|
|
FALLBACK_MODEL directly and never touch FREE_MODELS at all, so a
|
|
dead/rate-limited free-tier lineup can't add latency or fail a job.
|
|
"""
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1', 'model-free-2'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
monkeypatch.setattr(extractor, 'USE_FREE_MODELS_FIRST', False)
|
|
|
|
calls = []
|
|
|
|
def _fake_call(model, prompt, content):
|
|
calls.append(model)
|
|
return {'tripName': 'ok'}
|
|
|
|
monkeypatch.setattr(extractor, '_call_openrouter', _fake_call)
|
|
sentry_calls = []
|
|
monkeypatch.setattr(extractor.sentry_sdk, 'capture_message', lambda msg, level=None: sentry_calls.append(msg))
|
|
|
|
result = extract_with_openrouter('prompt')
|
|
|
|
assert result == {'tripName': 'ok'}
|
|
assert calls == ['model-paid'] # free models never attempted
|
|
assert sentry_calls == [] # no "free models exhausted" warning — none were tried
|