115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
"""Regresiones de selección de modelo en trabajos ARQ."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
|
|
from worker import tasks
|
|
|
|
|
|
class FakeSession:
|
|
def __init__(self, task: object) -> None:
|
|
self.task = task
|
|
|
|
def get(self, _model: object, _identifier: str, **_kwargs: object) -> object:
|
|
return self.task
|
|
|
|
def add(self, _value: object) -> None:
|
|
pass
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
|
|
class FakeRateLimiter:
|
|
async def acquire(self, *, estimated_tokens: int) -> None:
|
|
assert estimated_tokens >= 2000
|
|
|
|
|
|
class FakeService:
|
|
def __init__(self, model: str) -> None:
|
|
self.extractor = SimpleNamespace(model=model)
|
|
self.calls = 0
|
|
|
|
def extract_full_from_pdf(self, *_args: object, **_kwargs: object):
|
|
self.calls += 1
|
|
return {"numero": "F-1"}, {}
|
|
|
|
|
|
def test_process_document_uses_requested_model(monkeypatch, tmp_path) -> None:
|
|
task = SimpleNamespace(
|
|
request={
|
|
"model": "modelo-elegido",
|
|
"fields": ["numero"],
|
|
"structure_mode": "general",
|
|
"doc_types": [],
|
|
},
|
|
result={"processed": 0, "total": 1, "row_count": 0},
|
|
updated_at=None,
|
|
)
|
|
session = FakeSession(task)
|
|
|
|
@contextmanager
|
|
def fake_session_scope():
|
|
yield session
|
|
|
|
base_service = FakeService("modelo-global")
|
|
selected_service = FakeService("modelo-elegido")
|
|
selected_models: list[str] = []
|
|
|
|
def fake_service_for_model(_base: object, model: str):
|
|
selected_models.append(model)
|
|
return selected_service
|
|
|
|
monkeypatch.setattr(tasks, "session_scope", fake_session_scope)
|
|
monkeypatch.setattr(tasks, "get_rate_limiter", lambda: FakeRateLimiter())
|
|
monkeypatch.setattr(tasks, "_service_for_model", fake_service_for_model)
|
|
monkeypatch.setattr(tasks, "normalize_to_schema", lambda data, **_kwargs: {"normalized": data})
|
|
monkeypatch.setattr(tasks, "flatten_extraction_result", lambda *_args, **_kwargs: [])
|
|
monkeypatch.setattr(tasks, "flag_modified", lambda *_args, **_kwargs: None)
|
|
monkeypatch.setattr(tasks, "mark_task_completed_if_done", lambda *_args, **_kwargs: None)
|
|
|
|
queue_file = tmp_path / "factura.pdf"
|
|
queue_file.write_bytes(b"pdf")
|
|
ctx = {"service": base_service}
|
|
|
|
asyncio.run(tasks.process_document(
|
|
ctx,
|
|
task_id="task-1",
|
|
batch_id="batch-1",
|
|
queue_path=str(queue_file),
|
|
radicado="RAD-1",
|
|
filename="factura.pdf",
|
|
doc_type=None,
|
|
instruction=None,
|
|
tiff_page_range=None,
|
|
))
|
|
|
|
assert selected_models == ["modelo-elegido"]
|
|
assert selected_service.calls == 1
|
|
assert base_service.calls == 0
|
|
assert ctx["services_by_model"]["modelo-elegido"] is selected_service
|
|
|
|
|
|
def test_service_for_request_reuses_cached_model(monkeypatch) -> None:
|
|
base_service = FakeService("modelo-global")
|
|
selected_service = FakeService("modelo-elegido")
|
|
builds = 0
|
|
|
|
def fake_service_for_model(_base: object, _model: str):
|
|
nonlocal builds
|
|
builds += 1
|
|
return selected_service
|
|
|
|
monkeypatch.setattr(tasks, "_service_for_model", fake_service_for_model)
|
|
ctx = {"service": base_service}
|
|
request = {"model": "modelo-elegido"}
|
|
|
|
first = tasks._service_for_request(ctx, request)
|
|
second = tasks._service_for_request(ctx, request)
|
|
|
|
assert first is selected_service
|
|
assert second is selected_service
|
|
assert builds == 1
|