Testing and some tweaks

pyproject.toml

@@ -51,6 +51,8 @@ poethepoet = "^0.26.1"
 pytest = "^8.2.0"
 pytest-cov = "^5.0.0"
 pytest-asyncio = "^0.23.6"
+nest-asyncio = "^1.6.0"
+freezegun = "^1.5.1"
 
 [tool.poetry.extras]
 postgres = ["aiopg", "databases", "asyncpg", "sqlparse"]
@@ -79,6 +81,10 @@ sequence = [
     {cmd = "pytest -vv --cov-report xml --cov-report term --cov=flux tests/unit"},
 ]
 
+[tool.poe.tasks.unit.env]
+FLUX_TESTING = "1"
+
+
 [tool.poe.tasks.integration]
 sequence = [
     {cmd = "docker compose -f docker-compose.test.yml up -d --wait"},
@@ -90,6 +96,7 @@ ignore_fail = "return_non_zero"
 [tool.poe.tasks.integration.env]
 TEST_PG_CONNECTION_STRING = "postgresql+aiopg://postgres:postgres@localhost:55443"
 TEST_PG_MANAGEMENT_DB = "postgres"
+FLUX_TESTING = "1"
 
 
 [tool.poe.tasks.integration-ci]
@@ -101,6 +108,7 @@ ignore_fail = "return_non_zero"
 [tool.poe.tasks.integration-ci.env]
 TEST_PG_CONNECTION_STRING = "postgresql+aiopg://postgres:postgres@localhost:5432"
 TEST_PG_MANAGEMENT_DB = "postgres"
+FLUX_TESTING = "1"
 
 
 [build-system]

src/flux/cli.py

@@ -1,6 +1,6 @@
+import asyncio
 import datetime as dt
 import os
-from asyncio import run as asyncio_run
 from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
@@ -24,6 +24,16 @@
 NOT_APPLIED_STATUS = "Not Applied"
 
 
+def async_run(coro):
+    # Temp ugly workaround for testing until typer supports async
+    # See fastapi/typer#950
+    if os.environ.get("FLUX_TESTING"):
+        import nest_asyncio
+
+        nest_asyncio.apply()
+    asyncio.run(coro)
+
+
 @dataclass
 class _CliState:
     config: FluxConfig | None = None
@@ -167,7 +177,7 @@ def new(
     name: Annotated[str, typer.Argument(help="Migration name and default comment")],
     kind: MigrationKind = MigrationKind.python,
 ):
-    asyncio_run(_new(ctx=ctx, name=name, kind=kind))
+    async_run(_new(ctx=ctx, name=name, kind=kind))
 
 
 async def _print_apply_report(runner: FluxRunner, n: int | None):
@@ -249,7 +259,7 @@ def apply(
     ] = None,
     auto_approve: bool = False,
 ):
-    asyncio_run(
+    async_run(
         _apply(ctx, connection_uri=connection_uri, n=n, auto_approve=auto_approve)
     )
 
@@ -290,6 +300,6 @@ def rollback(
     ] = None,
     auto_approve: bool = False,
 ):
-    asyncio_run(
+    async_run(
         _rollback(ctx, connection_uri=connection_uri, n=n, auto_approve=auto_approve)
     )

src/flux/runner.py

@@ -165,7 +165,9 @@ async def apply_migrations(self, n: int | None = None):
 
         self.applied_migrations = await self.backend.get_applied_migrations()
 
-    def migrations_to_rollback(self, n: int | None = None):
+    def migrations_to_rollback(self, n: int | None = None) -> list[Migration]:
+        if n == 0:
+            return []
         applied_migrations = self.list_applied_migrations()
         migrations_to_rollback = (
             applied_migrations[-n:] if n is not None else applied_migrations
@@ -201,3 +203,23 @@ async def rollback_migrations(self, n: int | None = None):
                     await self._apply_post_apply_migrations()
 
         self.applied_migrations = await self.backend.get_applied_migrations()
+
+    async def rollback_migration(self, migration_id: str):
+        """
+        Rollback all migrations up to and including the given migration ID
+        """
+        applied_migrations = self.list_applied_migrations()
+        target_migration_index = next(
+            (
+                index
+                for index, migration in enumerate(applied_migrations)
+                if migration.id == migration_id
+            ),
+            None,
+        )
+        if target_migration_index is None:
+            raise ValueError(f"Migration {migration_id!r} has not been applied")
+
+        n = len(applied_migrations) - target_migration_index
+
+        await self.rollback_migrations(n=n)

tests/integration/postgres/conftest.py

@@ -28,21 +28,31 @@ async def test_database() -> AsyncGenerator[str, None]:
 
 
 @pytest.fixture
-async def postgres_backend(test_database):
+def database_uri(test_database) -> str:
+    return f"{TEST_PG_CONNECTION_STRING}/{test_database}"
+
+
+@pytest.fixture
+async def postgres_backend(database_uri: str):
     return FluxPostgresBackend(
-        database_url=f"{TEST_PG_CONNECTION_STRING}/{test_database}",
-        migrations_schema="_migrations",
+        database_url=database_uri,
         migrations_table="_flux_migrations",
     )
 
 
 @pytest.fixture
-def example_migrations_dir() -> Generator[str, None, None]:
+def example_project_dir() -> Generator[str, None, None]:
+    with TemporaryDirectory() as tempdir:
+        migrations_dir = os.path.join(tempdir, "migrations")
+        shutil.copytree(MIGRATIONS_1_DIR, migrations_dir)
+        yield tempdir
+
+
+@pytest.fixture
+def example_migrations_dir(example_project_dir: str) -> str:
     """
     Create a temporary copy of the migrations directory that can be freely
     modified by tests
     """
-    with TemporaryDirectory() as tempdir:
-        migrations_dir = os.path.join(tempdir, "migrations")
-        shutil.copytree(MIGRATIONS_1_DIR, migrations_dir)
-        yield migrations_dir
+    migrations_dir = os.path.join(example_project_dir, "migrations")
+    return migrations_dir