from __future__ import annotations import json import subprocess from datetime import UTC, datetime from decimal import Decimal from pathlib import Path from threading import Semaphore from types import SimpleNamespace import pytest from commercial_runtime_testkit import seed_meter from sqlalchemy import create_engine, select from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import StaticPool import app.models # noqa: F401 - 注册完整 metadata from app.core.config import get_settings from app.db.base_class import Base from app.models.commercial import UsageMeterEvent from app.models.commercial_runtime import CommercialRuntimeReservation from app.services.commercial_direct_operation import CommercialDirectOperationBridge from app.services.ocr import WORKER_JSON_PREFIX, OcrService from app.services.ocr_commercial import ( OcrCommercialAccessDenied, OcrCommercialObserver, OcrOperationContext, ) from app.services.ocr_worker_runtime import invoke_ocr_worker @pytest.fixture() def factory() -> sessionmaker[Session]: engine = create_engine( "sqlite+pysqlite:///:memory:", connect_args={"check_same_thread": False}, poolclass=StaticPool, ) Base.metadata.create_all(engine) result = sessionmaker(bind=engine, expire_on_commit=False) try: yield result finally: Base.metadata.drop_all(engine) engine.dispose() def _seed_ocr_meter( factory: sessionmaker[Session], *, preflight_quantity: Decimal = Decimal("5"), hard_limit: Decimal = Decimal("20"), ) -> None: with factory() as db: seed_meter( db, "tenant-a", datetime.now(UTC), basis="pages", tool_type="ocr", tool_name="paddle.worker", preflight_quantity=preflight_quantity, hard_limit=hard_limit, ) db.commit() def _context(suffix: str = "one") -> OcrOperationContext: return OcrOperationContext( tenant_id="tenant-a", operation_id=f"ocr-operation-{suffix}", run_id=f"ocr-run-{suffix}", ) def _settings() -> SimpleNamespace: return SimpleNamespace( ocr_language="ch", ocr_device="", ocr_text_detection_model="PP-OCRv5_mobile_det", ocr_text_recognition_model="PP-OCRv5_mobile_rec", ocr_timeout_seconds=10, ) def _invoke( *, observer: OcrCommercialObserver, context: OcrOperationContext, input_paths: list[Path], ) -> dict: return invoke_ocr_worker( settings=_settings(), python_bin="python", worker_path="worker.py", input_paths=input_paths, semaphore=Semaphore(1), parse_stdout=_parse_json, commercial_observer=observer, operation_context=context, ) def _parse_json(value: str) -> dict | None: try: payload = json.loads(value) except json.JSONDecodeError: return None return payload if isinstance(payload, dict) else None def test_ocr_worker_settles_exact_prepared_page_count_once( factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: _seed_ocr_meter(factory) observer = OcrCommercialObserver(CommercialDirectOperationBridge(factory)) inputs = [tmp_path / "page-1.png", tmp_path / "page-2.png"] monkeypatch.setattr( "app.services.ocr_worker_runtime.subprocess.run", lambda *args, **kwargs: subprocess.CompletedProcess( args=args[0], returncode=0, stdout='{"documents": []}', stderr="" ), ) payload = _invoke(observer=observer, context=_context(), input_paths=inputs) assert payload == {"documents": []} with factory() as db: reservation = db.scalars(select(CommercialRuntimeReservation)).one() usage = db.scalars(select(UsageMeterEvent)).one() assert reservation.status == "committed" assert Decimal(reservation.reserved_quantity) == Decimal("2") assert Decimal(reservation.actual_quantity or 0) == Decimal("2") assert Decimal(usage.quantity) == Decimal("2") @pytest.mark.parametrize( ("returncode", "stdout", "expected_error"), [ (3, "", "OCR 执行失败"), (0, "not-json", "JSON"), ], ) def test_ocr_worker_failure_after_send_still_records_real_pages( factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, tmp_path: Path, returncode: int, stdout: str, expected_error: str, ) -> None: _seed_ocr_meter(factory) observer = OcrCommercialObserver(CommercialDirectOperationBridge(factory)) monkeypatch.setattr( "app.services.ocr_worker_runtime.subprocess.run", lambda *args, **kwargs: subprocess.CompletedProcess( args=args[0], returncode=returncode, stdout=stdout, stderr="provider-error" ), ) with pytest.raises((RuntimeError, json.JSONDecodeError), match=expected_error): _invoke( observer=observer, context=_context(f"failure-{returncode}"), input_paths=[tmp_path / "page.png"], ) with factory() as db: usage = db.scalars(select(UsageMeterEvent)).one() assert Decimal(usage.quantity) == Decimal("1") def test_ocr_quota_is_checked_before_subprocess( factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: _seed_ocr_meter(factory, preflight_quantity=Decimal("1")) observer = OcrCommercialObserver(CommercialDirectOperationBridge(factory)) called = False def fail_if_called(*args, **kwargs): nonlocal called called = True raise AssertionError("额度拒绝后不应启动 OCR 子进程。") monkeypatch.setattr("app.services.ocr_worker_runtime.subprocess.run", fail_if_called) with pytest.raises(OcrCommercialAccessDenied, match="preflight_quantity"): _invoke( observer=observer, context=_context("over-limit"), input_paths=[tmp_path / "page-1.png", tmp_path / "page-2.png"], ) assert called is False with factory() as db: assert db.query(CommercialRuntimeReservation).count() == 0 assert db.query(UsageMeterEvent).count() == 0 def test_ocr_cache_hit_does_not_start_or_bill_worker_again( factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: _seed_ocr_meter(factory) calls = 0 def fake_run(command, **kwargs): nonlocal calls calls += 1 input_path = command[command.index("--input") + 1] payload = { "engine": "paddleocr_mobile", "model": "PP-OCRv5_mobile_rec", "documents": [ { "input_path": input_path, "text": "增值税发票 金额 100 元", "summary": "增值税发票", "line_count": 1, "page_count": 1, "lines": [{"text": "增值税发票 金额 100 元", "score": 0.98}], } ], } return subprocess.CompletedProcess( args=command, returncode=0, stdout=f"{WORKER_JSON_PREFIX}{json.dumps(payload, ensure_ascii=False)}\n", stderr="", ) monkeypatch.setenv("STORAGE_ROOT_DIR", str(tmp_path / "storage")) monkeypatch.setattr("app.services.ocr_worker_runtime.subprocess.run", fake_run) monkeypatch.setattr(OcrService, "_resolve_python_bin", lambda self: "python") monkeypatch.setattr(OcrService, "_resolve_worker_path", lambda self: "worker.py") get_settings.cache_clear() OcrService.clear_result_cache() content = b"same-real-image" try: with factory() as db: first = OcrService(db, operation_context=_context("cache-first")) second = OcrService(db, operation_context=_context("cache-second")) first.recognize_files([("invoice.png", content, "image/png")]) second.recognize_files([("renamed.png", content, "image/png")]) finally: OcrService.clear_result_cache() get_settings.cache_clear() assert calls == 1 with factory() as db: assert db.query(UsageMeterEvent).count() == 1