generated from john/python-template
This commit is contained in:
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy import select
|
||||
from sqlmodel import SQLModel
|
||||
|
||||
from transcription.db.migration import export_bundle
|
||||
from transcription.db.migration import import_bundle
|
||||
from transcription.db.migration import sqlite_url_from_path
|
||||
|
||||
# Register table metadata.
|
||||
from transcription.db import models as _models # noqa: F401
|
||||
|
||||
|
||||
def test_export_import_migration_round_trips_db_and_uploads(tmp_path):
|
||||
source_db_path = tmp_path / "source.db"
|
||||
target_db_path = tmp_path / "target.db"
|
||||
source_upload_dir = tmp_path / "source_uploads"
|
||||
target_upload_dir = tmp_path / "target_uploads"
|
||||
bundle_dir = tmp_path / "bundle"
|
||||
|
||||
source_db_url = sqlite_url_from_path(source_db_path)
|
||||
target_db_url = sqlite_url_from_path(target_db_path)
|
||||
|
||||
document_id = uuid4()
|
||||
job_id = uuid4()
|
||||
source_id = uuid4()
|
||||
job_source_id = uuid4()
|
||||
attempt_id = uuid4()
|
||||
filename = f"{source_id}.jpg"
|
||||
media_path = source_upload_dir / "documents" / str(document_id) / filename
|
||||
media_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
media_path.write_bytes(b"sample-image")
|
||||
|
||||
engine = create_engine(source_db_url)
|
||||
try:
|
||||
SQLModel.metadata.create_all(engine)
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
SQLModel.metadata.tables["document"].insert(),
|
||||
[{"id": document_id, "name": "Export doc"}],
|
||||
)
|
||||
connection.execute(
|
||||
SQLModel.metadata.tables["job"].insert(),
|
||||
[{"id": job_id, "document_id": document_id, "status": "queued"}],
|
||||
)
|
||||
connection.execute(
|
||||
SQLModel.metadata.tables["source"].insert(),
|
||||
[
|
||||
{
|
||||
"id": source_id,
|
||||
"document_id": document_id,
|
||||
"page_number": 1,
|
||||
"upload_name": "upload.jpg",
|
||||
"filename": filename,
|
||||
"file_path": str(media_path),
|
||||
"file_hash": "a" * 64,
|
||||
"file_size_bytes": len(b"sample-image"),
|
||||
}
|
||||
],
|
||||
)
|
||||
connection.execute(
|
||||
SQLModel.metadata.tables["job_source"].insert(),
|
||||
[{"id": job_source_id, "job_id": job_id, "source_id": source_id, "status": "pending"}],
|
||||
)
|
||||
connection.execute(
|
||||
SQLModel.metadata.tables["execution_attempt"].insert(),
|
||||
[
|
||||
{
|
||||
"id": attempt_id,
|
||||
"job_source_id": job_source_id,
|
||||
"job_id": job_id,
|
||||
"source_id": source_id,
|
||||
"attempt_number": 1,
|
||||
"status": "transcribed",
|
||||
"provider": "fixture",
|
||||
"model": "fixture-model",
|
||||
"transport_body": b"body",
|
||||
"raw_transcription": "hello",
|
||||
"started_at": datetime.now(UTC),
|
||||
"finished_at": datetime.now(UTC),
|
||||
"duration_ms": 10,
|
||||
}
|
||||
],
|
||||
)
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
export_bundle(source_db_url=source_db_url, source_upload_dir=source_upload_dir, bundle_dir=bundle_dir)
|
||||
import_bundle(target_db_url=target_db_url, target_upload_dir=target_upload_dir, bundle_dir=bundle_dir)
|
||||
|
||||
target_engine = create_engine(target_db_url)
|
||||
try:
|
||||
with target_engine.connect() as connection:
|
||||
source_row = connection.execute(
|
||||
select(SQLModel.metadata.tables["source"].c.file_path).where(
|
||||
SQLModel.metadata.tables["source"].c.id == source_id
|
||||
)
|
||||
).one()
|
||||
attempt_row = connection.execute(
|
||||
select(SQLModel.metadata.tables["execution_attempt"].c.transport_body).where(
|
||||
SQLModel.metadata.tables["execution_attempt"].c.id == attempt_id
|
||||
)
|
||||
).one()
|
||||
assert source_row[0] == f"documents/{document_id}/{filename}"
|
||||
assert attempt_row[0] == b"body"
|
||||
finally:
|
||||
target_engine.dispose()
|
||||
|
||||
copied_media_path = target_upload_dir / "documents" / str(document_id) / filename
|
||||
assert copied_media_path.read_bytes() == b"sample-image"
|
||||
Reference in New Issue
Block a user