moving to asyncengine

Author: sanderegg

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_manager.py

@@ -325,7 +325,7 @@ async def test_empty_pipeline_is_not_scheduled(
     initialized_app: FastAPI,
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
-    pipeline: Callable[..., CompPipelineAtDB],
+    create_pipeline: Callable[..., Awaitable[CompPipelineAtDB]],
     run_metadata: RunMetadataDict,
     sqlalchemy_async_engine: AsyncEngine,
     scheduler_rabbit_client_parser: mock.AsyncMock,
@@ -349,7 +349,7 @@ async def test_empty_pipeline_is_not_scheduled(
     scheduler_rabbit_client_parser.assert_not_called()
 
     # create the empty pipeline now
-    pipeline(project_id=f"{empty_project.uuid}")
+    await create_pipeline(project_id=f"{empty_project.uuid}")
 
     # creating a run with an empty pipeline is useless, check the scheduler is not kicking in
     with caplog.at_level(logging.WARNING):

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_scheduler_dask.py

@@ -1091,7 +1091,7 @@ async def test_broken_pipeline_configuration_is_not_scheduled_and_aborted(
     scheduler_api: BaseCompScheduler,
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
-    pipeline: Callable[..., CompPipelineAtDB],
+    create_pipeline: Callable[..., Awaitable[CompPipelineAtDB]],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
     sqlalchemy_async_engine: AsyncEngine,
@@ -1101,7 +1101,7 @@ async def test_broken_pipeline_configuration_is_not_scheduled_and_aborted(
     It shall be aborted and shown as such in the comp_runs db"""
     user = registered_user()
     sleepers_project = await project(user, workbench=fake_workbench_without_outputs)
-    pipeline(
+    await create_pipeline(
         project_id=f"{sleepers_project.uuid}",
         dag_adjacency_list=fake_workbench_adjacency,
     )

services/director-v2/tests/unit/with_dbs/conftest.py

@@ -6,8 +6,8 @@
 
 
 import datetime
-from collections.abc import Awaitable, Callable, Iterator
-from typing import Any, AsyncIterator, cast
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Any, cast
 from uuid import uuid4
 
 import arrow
@@ -40,20 +40,20 @@
 
 
 @pytest.fixture
-def pipeline(
-    postgres_db: sa.engine.Engine,
-) -> Iterator[Callable[..., CompPipelineAtDB]]:
+async def create_pipeline(
+    sqlalchemy_async_engine: AsyncEngine,
+) -> AsyncIterator[Callable[..., Awaitable[CompPipelineAtDB]]]:
     created_pipeline_ids: list[str] = []
 
-    def creator(**pipeline_kwargs) -> CompPipelineAtDB:
+    async def _(**pipeline_kwargs) -> CompPipelineAtDB:
         pipeline_config = {
             "project_id": f"{uuid4()}",
             "dag_adjacency_list": {},
             "state": StateType.NOT_STARTED,
         }
         pipeline_config.update(**pipeline_kwargs)
-        with postgres_db.begin() as conn:
-            result = conn.execute(
+        async with sqlalchemy_async_engine.begin() as conn:
+            result = await conn.execute(
                 comp_pipeline.insert()
                 .values(**pipeline_config)
                 .returning(sa.literal_column("*"))
@@ -64,11 +64,11 @@ def creator(**pipeline_kwargs) -> CompPipelineAtDB:
             created_pipeline_ids.append(f"{new_pipeline.project_id}")
             return new_pipeline
 
-    yield creator
+    yield _
 
     # cleanup
-    with postgres_db.connect() as conn:
-        conn.execute(
+    async with sqlalchemy_async_engine.connect() as conn:
+        await conn.execute(
             comp_pipeline.delete().where(
                 comp_pipeline.c.project_id.in_(created_pipeline_ids)
             )
@@ -81,7 +81,7 @@ async def create_tasks(
 ) -> AsyncIterator[Callable[..., Awaitable[list[CompTaskAtDB]]]]:
     created_task_ids: list[int] = []
 
-    async def creator(
+    async def _(
         user: dict[str, Any], project: ProjectAtDB, **overrides_kwargs
     ) -> list[CompTaskAtDB]:
         created_tasks: list[CompTaskAtDB] = []
@@ -144,7 +144,7 @@ async def creator(
             created_task_ids.extend([t.task_id for t in created_tasks if t.task_id])
         return created_tasks
 
-    yield creator
+    yield _
 
     # cleanup
     async with sqlalchemy_async_engine.connect() as conn:
@@ -224,29 +224,37 @@ async def _(
 
 
 @pytest.fixture
-def cluster(
-    postgres_db: sa.engine.Engine,
-) -> Iterator[Callable[..., Cluster]]:
+async def create_cluster(
+    sqlalchemy_async_engine: AsyncEngine,
+) -> AsyncIterator[Callable[..., Awaitable[Cluster]]]:
     created_cluster_ids: list[str] = []
 
-    def creator(user: dict[str, Any], **cluster_kwargs) -> Cluster:
+    async def _(user: dict[str, Any], **cluster_kwargs) -> Cluster:
+        assert "json_schema_extra" in Cluster.model_config
+        assert isinstance(Cluster.model_config["json_schema_extra"], dict)
+        assert isinstance(Cluster.model_config["json_schema_extra"]["examples"], list)
+        assert isinstance(
+            Cluster.model_config["json_schema_extra"]["examples"][1], dict
+        )
         cluster_config = Cluster.model_config["json_schema_extra"]["examples"][1]
         cluster_config["owner"] = user["primary_gid"]
         cluster_config.update(**cluster_kwargs)
         new_cluster = Cluster.model_validate(cluster_config)
         assert new_cluster
 
-        with postgres_db.connect() as conn:
+        async with sqlalchemy_async_engine.connect() as conn:
             # insert basic cluster
-            created_cluster = conn.execute(
-                sa.insert(clusters)
-                .values(to_clusters_db(new_cluster, only_update=False))
-                .returning(sa.literal_column("*"))
+            created_cluster = (
+                await conn.execute(
+                    sa.insert(clusters)
+                    .values(to_clusters_db(new_cluster, only_update=False))
+                    .returning(sa.literal_column("*"))
+                )
             ).one()
             created_cluster_ids.append(created_cluster.id)
             if "access_rights" in cluster_kwargs:
                 for gid, rights in cluster_kwargs["access_rights"].items():
-                    conn.execute(
+                    await conn.execute(
                         pg_insert(cluster_to_groups)
                         .values(
                             cluster_id=created_cluster.id,
@@ -259,7 +267,7 @@ def creator(user: dict[str, Any], **cluster_kwargs) -> Cluster:
                         )
                     )
             access_rights_in_db = {}
-            for row in conn.execute(
+            for row in await conn.execute(
                 sa.select(
                     cluster_to_groups.c.gid,
                     cluster_to_groups.c.read,
@@ -287,12 +295,11 @@ def creator(user: dict[str, Any], **cluster_kwargs) -> Cluster:
                 thumbnail=None,
             )
 
-    yield creator
+    yield _
 
     # cleanup
-    with postgres_db.connect() as conn:
-        conn.execute(
-            # pylint: disable=no-value-for-parameter
+    async with sqlalchemy_async_engine.connect() as conn:
+        await conn.execute(
             clusters.delete().where(clusters.c.id.in_(created_cluster_ids))
         )
 
@@ -301,8 +308,8 @@ def creator(user: dict[str, Any], **cluster_kwargs) -> Cluster:
 async def publish_project(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
-    pipeline: Callable[..., CompPipelineAtDB],
-    create_tasks: Callable[..., list[CompTaskAtDB]],
+    create_pipeline: Callable[..., Awaitable[CompPipelineAtDB]],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> Callable[[], Awaitable[PublishedProject]]:
@@ -312,11 +319,11 @@ async def _() -> PublishedProject:
         created_project = await project(user, workbench=fake_workbench_without_outputs)
         return PublishedProject(
             project=created_project,
-            pipeline=pipeline(
+            pipeline=await create_pipeline(
                 project_id=f"{created_project.uuid}",
                 dag_adjacency_list=fake_workbench_adjacency,
             ),
-            tasks=create_tasks(
+            tasks=await create_tasks(
                 user=user, project=created_project, state=StateType.PUBLISHED
             ),
         )
@@ -335,9 +342,9 @@ async def published_project(
 async def running_project(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
-    pipeline: Callable[..., CompPipelineAtDB],
-    create_tasks: Callable[..., list[CompTaskAtDB]],
-    create_comp_run: Callable[..., CompRunsAtDB],
+    create_pipeline: Callable[..., Awaitable[CompPipelineAtDB]],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
+    create_comp_run: Callable[..., Awaitable[CompRunsAtDB]],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> RunningProject:
@@ -346,18 +353,18 @@ async def running_project(
     now_time = arrow.utcnow().datetime
     return RunningProject(
         project=created_project,
-        pipeline=pipeline(
+        pipeline=await create_pipeline(
             project_id=f"{created_project.uuid}",
             dag_adjacency_list=fake_workbench_adjacency,
         ),
-        tasks=create_tasks(
+        tasks=await create_tasks(
             user=user,
             project=created_project,
             state=StateType.RUNNING,
             progress=0.0,
             start=now_time,
         ),
-        runs=create_comp_run(
+        runs=await create_comp_run(
             user=user,
             project=created_project,
             started=now_time,
@@ -371,9 +378,9 @@ async def running_project(
 async def running_project_mark_for_cancellation(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
-    pipeline: Callable[..., CompPipelineAtDB],
-    create_tasks: Callable[..., list[CompTaskAtDB]],
-    create_comp_run: Callable[..., CompRunsAtDB],
+    create_pipeline: Callable[..., Awaitable[CompPipelineAtDB]],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
+    create_comp_run: Callable[..., Awaitable[CompRunsAtDB]],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> RunningProject:
@@ -382,18 +389,18 @@ async def running_project_mark_for_cancellation(
     now_time = arrow.utcnow().datetime
     return RunningProject(
         project=created_project,
-        pipeline=pipeline(
+        pipeline=await create_pipeline(
             project_id=f"{created_project.uuid}",
             dag_adjacency_list=fake_workbench_adjacency,
         ),
-        tasks=create_tasks(
+        tasks=await create_tasks(
             user=user,
             project=created_project,
             state=StateType.RUNNING,
             progress=0.0,
             start=now_time,
         ),
-        runs=create_comp_run(
+        runs=await create_comp_run(
             user=user,
             project=created_project,
             result=StateType.RUNNING,