Added test for Async Engine migrations

python-ellar · eadwinCode · Jan 16, 2024 · Jan 14, 2024 · Jan 16, 2024 · Jan 16, 2024
commit 8e20ce5e72229fe719762ffac22447becb3d7a9e
diff --git a/tests/test_migrations/samples/default_async.py b/tests/test_migrations/samples/default_async.py
@@ -0,0 +1,46 @@
+#!/bin/env python
+import ellar_cli.click as click
+from ellar.app import AppFactory, current_injector
+from ellar.common.utils.importer import get_main_directory_by_stack
+from ellar_cli.main import create_ellar_cli
+from models import User
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ellar_sqlalchemy import EllarSQLAlchemyModule
+
+
+def bootstrap():
+    path = get_main_directory_by_stack(
+        "__main__/__parent__/__parent__/dumbs/default_async", stack_level=1
+    )
+    application = AppFactory.create_app(
+        modules=[
+            EllarSQLAlchemyModule.setup(
+                databases="sqlite+aiosqlite:///app.db",
+                migration_options={"context_configure": {"compare_types": False}},
+                root_path=str(path),
+            )
+        ]
+    )
+    return application
+
+
+cli = create_ellar_cli("default_async:bootstrap")
+
+
+@cli.command()
+@click.run_as_async
+async def add_user():
+    session = current_injector.get(AsyncSession)
+    user = User(name="default App Ellar")
+    session.add(user)
+
+    await session.commit()
+    await session.refresh(user)
+    await session.close()
+
+    click.echo(f"<User name={user.name} id={user.id}>")
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/tests/test_migrations/samples/multiple_database_async.py b/tests/test_migrations/samples/multiple_database_async.py
@@ -0,0 +1,57 @@
+#!/bin/env python
+import ellar_cli.click as click
+from ellar.app import AppFactory, current_injector
+from ellar.common.utils.importer import get_main_directory_by_stack
+from ellar_cli.main import create_ellar_cli
+from models import Group, User
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ellar_sqlalchemy import EllarSQLAlchemyModule
+
+
+def bootstrap():
+    path = get_main_directory_by_stack(
+        "__main__/__parent__/__parent__/dumbs/multiple_async", stack_level=1
+    )
+    application = AppFactory.create_app(
+        modules=[
+            EllarSQLAlchemyModule.setup(
+                databases={
+                    "default": "sqlite+aiosqlite:///app.db",
+                    "db1": "sqlite+aiosqlite:///app2.db",
+                },
+                migration_options={"context_configure": {"compare_types": False}},
+                root_path=str(path),
+            )
+        ]
+    )
+    return application
+
+
+cli = create_ellar_cli("multiple_database_async:bootstrap")
+
+
+@cli.command()
+@click.run_as_async
+async def add_user():
+    session = current_injector.get(AsyncSession)
+    user = User(name="Multiple Database App Ellar")
+    group = Group(name="group")
+
+    session.add(user)
+    session.add(group)
+
+    await session.commit()
+
+    await session.refresh(user)
+    await session.refresh(group)
+
+    await session.close()
+
+    click.echo(
+        f"<User name={user.name} id={user.id} and Group name={group.name} id={group.id}>"
+    )
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/tests/test_migrations/test_migrations_commands.py b/tests/test_migrations/test_migrations_commands.py
@@ -132,3 +132,33 @@ def test_other_alembic_commands():
     result = run_command("default.py db downgrade")
     assert result.returncode == 1
     assert b"Relative revision -1 didn't produce 1 migrations" in result.stderr
+
+
+@clean_directory("default_async")
+def test_migrate_upgrade_async():
+    result = run_command("default_async.py db init")
+    assert result.returncode == 0
+    assert (
+        b"tests/dumbs/default_async/migrations/alembic.ini' before proceeding."
+        in result.stdout
+    )
+
+    result = run_command("default_async.py db check")
+    assert result.returncode == 1
+
+    result = run_command("default_async.py db migrate")
+    assert result.returncode == 0
+
+    result = run_command("default_async.py db check")
+    assert result.returncode == 1
+
+    result = run_command("default_async.py db upgrade")
+    assert result.returncode == 0
+
+    result = run_command("default_async.py db check")
+    assert result.returncode == 0
+    assert result.stdout == b"No new upgrade operations detected.\n"
+
+    result = run_command("default_async.py add-user")
+    assert result.returncode == 0
+    assert result.stdout == b"<User name=default App Ellar id=1>\n"
diff --git a/tests/test_migrations/test_multiple_database.py b/tests/test_migrations/test_multiple_database.py
@@ -54,3 +54,37 @@ def test_migrate_upgrade_multiple_database_with_model_changes():
                 b"Detected type change from VARCHAR(length=256) to String(length=128)"
                 in result.stderr
             )
+
+
+@clean_directory("multiple_async")
+def test_migrate_upgrade_for_multiple_database_async():
+    with set_env_variable("multiple_db", "true"):
+        result = run_command("multiple_database_async.py db init")
+        assert result.returncode == 0
+        assert (
+            b"tests/dumbs/multiple_async/migrations/alembic.ini' before proceeding."
+            in result.stdout
+        )
+
+        result = run_command("multiple_database_async.py db check")
+        assert result.returncode == 1
+
+        result = run_command("multiple_database_async.py db migrate")
+        assert result.returncode == 0
+
+        result = run_command("multiple_database_async.py db check")
+        assert result.returncode == 1
+
+        result = run_command("multiple_database_async.py db upgrade")
+        assert result.returncode == 0
+
+        result = run_command("multiple_database_async.py db check")
+        assert result.returncode == 0
+        assert result.stdout == b"No new upgrade operations detected.\n"
+
+        result = run_command("multiple_database_async.py add-user")
+        assert result.returncode == 0
+        assert (
+            result.stdout
+            == b"<User name=Multiple Database App Ellar id=1 and Group name=group id=1>\n"
+        )