224 lines
11 KiB
Python
224 lines
11 KiB
Python
import json
|
||
from io import BytesIO
|
||
from datetime import date, datetime, timedelta, timezone
|
||
from types import SimpleNamespace
|
||
from unittest.mock import AsyncMock, Mock
|
||
from uuid import uuid4
|
||
import pytest
|
||
import pytest_asyncio
|
||
from fastapi import UploadFile
|
||
from sqlalchemy import select, func
|
||
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
|
||
from src.models import (RecognitionJob, StoredFile, FinancialStatementMeta,
|
||
FinancialStatementLine, FinancialStatementValue)
|
||
from src.models.base import Base
|
||
from src.schemas.enums import RecognitionStatus, RecognitionPipeline, FileStatus, DocumentType, FileType, StatementType, FslRowKind
|
||
from src.schemas.pdfworker import PDFWorkerResult
|
||
from src.services import pdfworker_results
|
||
from src.services.file_service import FileService
|
||
from src.repositories.file_repository import FileRepository
|
||
from src.repositories.recognition_job_repo import RecognitionJobRepository
|
||
from src.core import clients
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def db():
|
||
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
|
||
tables = [StoredFile.__table__, RecognitionJob.__table__, FinancialStatementMeta.__table__,
|
||
FinancialStatementLine.__table__, FinancialStatementValue.__table__]
|
||
async with engine.begin() as connection:
|
||
await connection.run_sync(lambda conn: Base.metadata.create_all(conn, tables=tables))
|
||
factory = async_sessionmaker(engine, expire_on_commit=False)
|
||
ids = SimpleNamespace(task=uuid4(), file=uuid4(), legal=uuid4())
|
||
async with factory() as session:
|
||
session.add(StoredFile(id=ids.file, legal_entity_id=ids.legal, file_type=FileType.PDF,
|
||
document_type=DocumentType.OFR,
|
||
original_name='report.pdf', s3_key='folder/nested/report.pdf',
|
||
size_bytes=123, period_year=2026, current_recognition_job_id=ids.task))
|
||
session.add(RecognitionJob(id=ids.task, file_id=ids.file, legal_entity_id=ids.legal,
|
||
pipeline=RecognitionPipeline.IMAGE, status=RecognitionStatus.QUEUED, attempt=1))
|
||
await session.commit()
|
||
yield factory, ids
|
||
await engine.dispose()
|
||
|
||
|
||
def ready_event(ids, status='ready'):
|
||
return PDFWorkerResult(task_id=ids.task, file_id=ids.file, status=status,
|
||
result_s3_key='results/full/path.json' if status == 'ready' else None,
|
||
error_code=None if status == 'ready' else 'storage_unavailable', finished_at=datetime.now(timezone.utc))
|
||
|
||
|
||
def setup_storage(ids, monkeypatch):
|
||
payload = dict(task_id=str(ids.task), job_id=str(ids.task), file_id=str(ids.file), files=[{}])
|
||
response = SimpleNamespace(read=AsyncMock(return_value=json.dumps(payload).encode()), release=Mock())
|
||
storage = SimpleNamespace(get_object=AsyncMock(return_value=response))
|
||
meta_id, line_id = uuid4(), uuid4()
|
||
converter = Mock(return_value=dict(
|
||
meta=[dict(id=meta_id, stored_file_id=ids.file, statement_type=StatementType.BALANCE,
|
||
inn='123', inn_verified=True, units='руб', units_verified=True,
|
||
report_date=date(2026,1,1), report_date_verified=True)],
|
||
rows=[dict(id=line_id, recognition_job_id=ids.task, meta_id=meta_id,
|
||
legal_entity_id=ids.legal, line_code='1600', line_code_recognized=True,
|
||
line_name='Итого', line_name_recognized=True, row_kind=FslRowKind.ROW, ordinal=1)],
|
||
values=[dict(id=uuid4(), line_id=line_id, amount=100, string_value='100', column_index=1, is_recognized=True)],
|
||
skipped_files=[]))
|
||
monkeypatch.setattr(pdfworker_results, 'json_to_db_rows', converter)
|
||
return storage, response, converter
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_ready_commits_once(db, monkeypatch):
|
||
factory, ids = db
|
||
storage, response, converter = setup_storage(ids, monkeypatch)
|
||
event = ready_event(ids)
|
||
await pdfworker_results.apply_pdfworker_result(event, factory, storage)
|
||
await pdfworker_results.apply_pdfworker_result(event, factory, storage)
|
||
storage.get_object.assert_awaited_once_with('results/full/path.json')
|
||
response.release.assert_called_once()
|
||
assert converter.call_args.args[0]['legal_entity_id'] == str(ids.legal)
|
||
async with factory() as session:
|
||
assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.DONE
|
||
stored_file = await session.get(StoredFile, ids.file)
|
||
assert stored_file.status == FileStatus.CHECK
|
||
assert stored_file.document_type == DocumentType.OFR
|
||
for model in (FinancialStatementMeta, FinancialStatementLine, FinancialStatementValue):
|
||
assert await session.scalar(select(func.count()).select_from(model)) == 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_upload_accepts_new_extension_and_saves_ofr_type(db):
|
||
factory, ids = db
|
||
storage = Mock(put_object=AsyncMock())
|
||
async with factory() as session:
|
||
service = FileService(storage, session, FileRepository(session),
|
||
RecognitionJobRepository(session), None)
|
||
uploaded = UploadFile(filename='report.docx', file=BytesIO(b'data'), size=4)
|
||
created = await service.save_file(uploaded, ids.legal, 'Q1', 2026)
|
||
stored = await session.get(StoredFile, created['file']['id'])
|
||
job = await session.get(RecognitionJob, created['rec_job']['id'])
|
||
assert stored.file_type == FileType.FILE
|
||
assert stored.document_type == DocumentType.OFR
|
||
assert job.file_id == stored.id
|
||
assert service._get_transformed_file_job_info(stored)['file_type'] == 'DOCX'
|
||
storage.put_object.assert_awaited_once_with(f'{stored.id}.docx', b'data')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize('status', ['error', 'timeout'])
|
||
async def test_error_without_minio(db, status):
|
||
factory, ids = db
|
||
data = ready_event(ids, 'error').model_dump()
|
||
if status == 'timeout':
|
||
data.update(status='timeout', error_code='deadline_exceeded')
|
||
event = PDFWorkerResult(**data)
|
||
storage = Mock(get_object=AsyncMock())
|
||
await pdfworker_results.apply_pdfworker_result(event, factory, storage)
|
||
await pdfworker_results.apply_pdfworker_result(event, factory, storage)
|
||
storage.get_object.assert_not_awaited()
|
||
async with factory() as session:
|
||
job = await session.get(RecognitionJob, ids.task)
|
||
assert job.status == RecognitionStatus.ERROR
|
||
assert event.error_code in job.error_message
|
||
assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_invalid_json_rolls_back(db, monkeypatch):
|
||
factory, ids = db
|
||
storage, response, _ = setup_storage(ids, monkeypatch)
|
||
response.read.return_value = b'{"files": []}'
|
||
with pytest.raises(ValueError):
|
||
await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
|
||
response.release.assert_called_once()
|
||
async with factory() as session:
|
||
assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED
|
||
assert await session.scalar(select(func.count()).select_from(FinancialStatementMeta)) == 0
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_pollworker_recovers_same_command(db, monkeypatch):
|
||
factory, ids = db
|
||
monkeypatch.setattr(clients, 'async_session_maker', factory)
|
||
worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: SimpleNamespace()))
|
||
first, = await worker.claim_pending_rows()
|
||
assert first.source_s3_key == 'folder/nested/report.pdf'
|
||
assert first.task_id == ids.task
|
||
assert not await worker.claim_pending_rows()
|
||
async with factory() as session:
|
||
job = await session.get(RecognitionJob, ids.task)
|
||
job.started_at = datetime.now(timezone.utc) - timedelta(minutes=2)
|
||
await session.commit()
|
||
retry, = await worker.claim_pending_rows()
|
||
assert first.model_dump_json() == retry.model_dump_json()
|
||
await worker.mark_published(first.task_id)
|
||
assert not await worker.claim_pending_rows()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_cannot_requeue_finished_job(db, monkeypatch):
|
||
factory, ids = db
|
||
monkeypatch.setattr(clients, 'async_session_maker', factory)
|
||
await pdfworker_results.apply_pdfworker_result(ready_event(ids, 'error'), factory, Mock())
|
||
worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None))
|
||
await worker.mark_queued(ids.task)
|
||
async with factory() as session:
|
||
assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.ERROR
|
||
assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_invalid_stored_key_does_not_block_queue(db, monkeypatch):
|
||
factory, ids = db
|
||
monkeypatch.setattr(clients, 'async_session_maker', factory)
|
||
async with factory() as session:
|
||
file = await session.get(StoredFile, ids.file)
|
||
file.s3_key = None
|
||
await session.commit()
|
||
worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None))
|
||
assert await worker.claim_pending_rows() == []
|
||
async with factory() as session:
|
||
job = await session.get(RecognitionJob, ids.task)
|
||
assert job.status == RecognitionStatus.ERROR
|
||
assert 'source_s3_key' in job.error_message
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_wrong_file_id_does_not_change_task(db):
|
||
factory, ids = db
|
||
data = ready_event(ids, 'error').model_dump()
|
||
data['file_id'] = uuid4()
|
||
with pytest.raises(ValueError, match='file_id'):
|
||
await pdfworker_results.apply_pdfworker_result(PDFWorkerResult(**data), factory, Mock())
|
||
async with factory() as session:
|
||
assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatches_flagged_document_type(db, monkeypatch):
|
||
factory, ids = db
|
||
storage, _, _ = setup_storage(ids, monkeypatch)
|
||
async with factory() as session:
|
||
stored = await session.get(StoredFile, ids.file)
|
||
stored.document_type = DocumentType.OSV
|
||
await session.commit()
|
||
handler = AsyncMock()
|
||
monkeypatch.setitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR, handler)
|
||
await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
|
||
handler.assert_awaited_once()
|
||
async with factory() as session:
|
||
job = await session.get(RecognitionJob, ids.task)
|
||
assert job.status == RecognitionStatus.DONE
|
||
assert (await session.get(StoredFile, ids.file)).document_type == DocumentType.OFR
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_missing_handler_does_not_complete_job(db, monkeypatch):
|
||
factory, ids = db
|
||
monkeypatch.delitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR)
|
||
storage = Mock(get_object=AsyncMock())
|
||
with pytest.raises(ValueError, match='handler'):
|
||
await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
|
||
storage.get_object.assert_not_awaited()
|
||
async with factory() as session:
|
||
assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED
|