133 lines
4.8 KiB
Python
133 lines
4.8 KiB
Python
import pytest
|
|
from core import extractor
|
|
from core.extractor import extract_with_openrouter, ExtractorError, RateLimitError, _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_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_extract_with_openrouter_tries_next_model_on_rate_limit(monkeypatch):
|
|
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 RateLimitError('rate limited')
|
|
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']
|
|
|
|
|
|
def test_extract_with_openrouter_falls_back_to_paid_model_and_logs_sentry(monkeypatch):
|
|
monkeypatch.setattr(extractor, 'FREE_MODELS', ['model-free-1'])
|
|
monkeypatch.setattr(extractor, 'FALLBACK_MODEL', 'model-paid')
|
|
|
|
def _fake_call(model, prompt, content):
|
|
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 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')
|