from __future__ import annotations

import importlib
import unittest

from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import create_engine, inspect


class PostboxMigrationTests(unittest.TestCase):
    def test_baseline_creates_and_drops_owned_tables(self) -> None:
        migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "c7d2e5f8a1b4_v010_postbox_baseline"
        )
        engine = create_engine("sqlite:///:memory:")
        try:
            with engine.begin() as connection:
                operations = Operations(MigrationContext.configure(connection))
                original = migration.op
                migration.op = operations
                try:
                    migration.upgrade()
                    tables = set(inspect(connection).get_table_names())
                    self.assertIn("postboxes", tables)
                    self.assertIn("postbox_messages", tables)
                    self.assertIn("postbox_deliveries", tables)
                    self.assertIn("postbox_access_events", tables)
                    message_columns = {
                        column["name"]
                        for column in inspect(connection).get_columns(
                            "postbox_messages"
                        )
                    }
                    self.assertTrue(
                        {
                            "ciphertext_ref",
                            "signed_manifest_ref",
                            "wrapped_keys",
                            "key_epoch",
                            "expires_at",
                            "withdrawn_at",
                        }.issubset(message_columns)
                    )
                    migration.downgrade()
                    self.assertFalse(
                        {
                            table
                            for table in inspect(connection).get_table_names()
                            if table.startswith("postbox")
                        }
                    )
                finally:
                    migration.op = original
        finally:
            engine.dispose()


if __name__ == "__main__":
    unittest.main()