generated from john/python-template
110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from transcription.db.models import Document
|
|
from transcription.db.models import Job
|
|
from transcription.db.models import JobSource
|
|
from transcription.db.models import JobSourceStatus
|
|
from transcription.db.models import JobStatus
|
|
from transcription.db.models import Source
|
|
from transcription.services.documents import DocumentService
|
|
from transcription.services.jobs import JobService
|
|
|
|
|
|
class TestJobService:
|
|
@pytest.mark.asyncio
|
|
async def test_create_and_read_job(self, job_service: JobService, document_service: DocumentService):
|
|
document = Document(id=uuid4(), name="test-bundle")
|
|
await document_service.create_document(document=document)
|
|
|
|
job = Job(document_id=document.id)
|
|
await job_service.create_job(job=job)
|
|
|
|
fetched = await job_service.read_job(job_id=job.id)
|
|
assert fetched.id == job.id
|
|
assert fetched.document is not None
|
|
assert fetched.document.id == document.id
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_job_state_updates_status_and_retry(
|
|
self,
|
|
job_service: JobService,
|
|
document_service: DocumentService,
|
|
):
|
|
document = Document(id=uuid4(), name="test-bundle")
|
|
await document_service.create_document(document=document)
|
|
|
|
job = Job(document_id=document.id)
|
|
await job_service.create_job(job=job)
|
|
|
|
updated = await job_service.update_job_state(
|
|
job_id=job.id,
|
|
status=JobStatus.PROCESSING,
|
|
retry_count_increment=1,
|
|
)
|
|
|
|
assert updated.status == JobStatus.PROCESSING
|
|
assert updated.retry_count == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_jobs_by_status(self, job_service: JobService, document_service: DocumentService):
|
|
document = Document(id=uuid4(), name="query-doc")
|
|
await document_service.create_document(document=document)
|
|
|
|
await job_service.create_job(job=Job(document_id=document.id, status=JobStatus.PROCESSING))
|
|
await job_service.create_job(job=Job(document_id=document.id, status=JobStatus.QUEUED))
|
|
|
|
result = await job_service.query_jobs(status=JobStatus.PROCESSING)
|
|
assert len(result) == 1
|
|
assert result[0].status == JobStatus.PROCESSING
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_jobs_by_source_filename(self, job_service: JobService, document_service: DocumentService):
|
|
document = Document(id=uuid4(), name="source-doc")
|
|
await document_service.create_document(document=document)
|
|
|
|
job = Job(document_id=document.id)
|
|
await job_service.create_job(job=job)
|
|
|
|
async with job_service._session_scope() as session:
|
|
source = Source(
|
|
document_id=document.id,
|
|
page_number=1,
|
|
upload_name="letter.jpg",
|
|
filename="stored-letter.jpg",
|
|
file_path="/uploads/stored-letter.jpg",
|
|
)
|
|
session.add(source)
|
|
await session.flush()
|
|
session.add(
|
|
JobSource(
|
|
job_id=job.id,
|
|
source_id=source.id,
|
|
status=JobSourceStatus.PENDING,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
result = await job_service.query_jobs(filename="stored-letter.jpg")
|
|
assert len(result) == 1
|
|
assert result[0].id == job.id
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_next_queued_job_orders_by_created_date(
|
|
self,
|
|
job_service: JobService,
|
|
document_service: DocumentService,
|
|
):
|
|
document = Document(id=uuid4(), name="ordered-doc")
|
|
await document_service.create_document(document=document)
|
|
|
|
first = Job(document_id=document.id, status=JobStatus.QUEUED)
|
|
second = Job(document_id=document.id, status=JobStatus.QUEUED)
|
|
await job_service.create_job(job=first)
|
|
await job_service.create_job(job=second)
|
|
|
|
next_job = await job_service.read_next_queued_job()
|
|
assert next_job is not None
|
|
assert next_job.id == first.id
|