from __future__ import annotations import tempfile import unittest from pathlib import Path from alembic.runtime.migration import MigrationContext from sqlalchemy import create_engine, inspect from govoplan_core.db.migrations import migrate_database from govoplan_dataflow.backend.manifest import get_manifest class DataflowMigrationTests(unittest.TestCase): def test_baseline_creates_dataflow_tables_and_head(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-dataflow-migration-") as directory: url = f"sqlite:///{Path(directory) / 'dataflow.db'}" migrate_database( database_url=url, enabled_modules=("dataflow",), manifest_factories=(get_manifest,), ) engine = create_engine(url) try: with engine.connect() as connection: self.assertIn( "d4f7a1c8e2b0", set(MigrationContext.configure(connection).get_current_heads()), ) self.assertEqual( { "dataflow_pipelines", "dataflow_pipeline_revisions", "dataflow_runs", }, { name for name in inspect(connection).get_table_names() if name.startswith("dataflow_") }, ) finally: engine.dispose() if __name__ == "__main__": unittest.main()