moving to asyncengine

Author: sanderegg

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

@@ -7,7 +7,7 @@
 
 import datetime
 from collections.abc import Awaitable, Callable, Iterator
-from typing import Any, cast
+from typing import Any, AsyncIterator, cast
 from uuid import uuid4
 
 import arrow
@@ -36,6 +36,7 @@
 from simcore_service_director_v2.utils.dask import generate_dask_job_id
 from simcore_service_director_v2.utils.db import to_clusters_db
 from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 
 @pytest.fixture
@@ -75,12 +76,12 @@ def creator(**pipeline_kwargs) -> CompPipelineAtDB:
 
 
 @pytest.fixture
-def tasks(
-    postgres_db: sa.engine.Engine,
-) -> Iterator[Callable[..., list[CompTaskAtDB]]]:
+async def create_tasks(
+    sqlalchemy_async_engine: AsyncEngine,
+) -> AsyncIterator[Callable[..., Awaitable[list[CompTaskAtDB]]]]:
     created_task_ids: list[int] = []
 
-    def creator(
+    async def creator(
         user: dict[str, Any], project: ProjectAtDB, **overrides_kwargs
     ) -> list[CompTaskAtDB]:
         created_tasks: list[CompTaskAtDB] = []
@@ -132,8 +133,8 @@ def creator(
                 ),
             }
             task_config.update(**overrides_kwargs)
-            with postgres_db.connect() as conn:
-                result = conn.execute(
+            async with sqlalchemy_async_engine.connect() as conn:
+                result = await conn.execute(
                     comp_tasks.insert()
                     .values(**task_config)
                     .returning(sa.literal_column("*"))
@@ -146,8 +147,8 @@ def creator(
     yield creator
 
     # cleanup
-    with postgres_db.connect() as conn:
-        conn.execute(
+    async with sqlalchemy_async_engine.connect() as conn:
+        await conn.execute(
             comp_tasks.delete().where(comp_tasks.c.task_id.in_(created_task_ids))
         )
 
@@ -186,12 +187,12 @@ def run_metadata(
 
 
 @pytest.fixture
-def runs(
-    postgres_db: sa.engine.Engine, run_metadata: RunMetadataDict
-) -> Iterator[Callable[..., CompRunsAtDB]]:
+async def create_comp_run(
+    sqlalchemy_async_engine: AsyncEngine, run_metadata: RunMetadataDict
+) -> AsyncIterator[Callable[..., Awaitable[CompRunsAtDB]]]:
     created_run_ids: list[int] = []
 
-    def creator(
+    async def _(
         user: dict[str, Any], project: ProjectAtDB, **run_kwargs
     ) -> CompRunsAtDB:
         run_config = {
@@ -203,8 +204,8 @@ def creator(
             "use_on_demand_clusters": False,
         }
         run_config.update(**run_kwargs)
-        with postgres_db.connect() as conn:
-            result = conn.execute(
+        async with sqlalchemy_async_engine.connect() as conn:
+            result = await conn.execute(
                 comp_runs.insert()
                 .values(**jsonable_encoder(run_config))
                 .returning(sa.literal_column("*"))
@@ -213,11 +214,13 @@ def creator(
             created_run_ids.append(new_run.run_id)
             return new_run
 
-    yield creator
+    yield _
 
     # cleanup
-    with postgres_db.connect() as conn:
-        conn.execute(comp_runs.delete().where(comp_runs.c.run_id.in_(created_run_ids)))
+    async with sqlalchemy_async_engine.connect() as conn:
+        await conn.execute(
+            comp_runs.delete().where(comp_runs.c.run_id.in_(created_run_ids))
+        )
 
 
 @pytest.fixture
@@ -299,7 +302,7 @@ async def publish_project(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
+    create_tasks: Callable[..., list[CompTaskAtDB]],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> Callable[[], Awaitable[PublishedProject]]:
@@ -313,7 +316,9 @@ async def _() -> PublishedProject:
                 project_id=f"{created_project.uuid}",
                 dag_adjacency_list=fake_workbench_adjacency,
             ),
-            tasks=tasks(user=user, project=created_project, state=StateType.PUBLISHED),
+            tasks=create_tasks(
+                user=user, project=created_project, state=StateType.PUBLISHED
+            ),
         )
 
     return _
@@ -331,8 +336,8 @@ async def running_project(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
-    runs: Callable[..., CompRunsAtDB],
+    create_tasks: Callable[..., list[CompTaskAtDB]],
+    create_comp_run: Callable[..., CompRunsAtDB],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> RunningProject:
@@ -345,14 +350,14 @@ async def running_project(
             project_id=f"{created_project.uuid}",
             dag_adjacency_list=fake_workbench_adjacency,
         ),
-        tasks=tasks(
+        tasks=create_tasks(
             user=user,
             project=created_project,
             state=StateType.RUNNING,
             progress=0.0,
             start=now_time,
         ),
-        runs=runs(
+        runs=create_comp_run(
             user=user,
             project=created_project,
             started=now_time,
@@ -367,8 +372,8 @@ async def running_project_mark_for_cancellation(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
-    runs: Callable[..., CompRunsAtDB],
+    create_tasks: Callable[..., list[CompTaskAtDB]],
+    create_comp_run: Callable[..., CompRunsAtDB],
     fake_workbench_without_outputs: dict[str, Any],
     fake_workbench_adjacency: dict[str, Any],
 ) -> RunningProject:
@@ -381,14 +386,14 @@ async def running_project_mark_for_cancellation(
             project_id=f"{created_project.uuid}",
             dag_adjacency_list=fake_workbench_adjacency,
         ),
-        tasks=tasks(
+        tasks=create_tasks(
             user=user,
             project=created_project,
             state=StateType.RUNNING,
             progress=0.0,
             start=now_time,
         ),
-        runs=runs(
+        runs=create_comp_run(
             user=user,
             project=created_project,
             result=StateType.RUNNING,

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

@@ -918,7 +918,7 @@ async def test_get_computation_from_not_started_computation_task(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
     async_client: httpx.AsyncClient,
 ):
     user = registered_user()
@@ -935,7 +935,7 @@ async def test_get_computation_from_not_started_computation_task(
     assert response.status_code == status.HTTP_409_CONFLICT, response.text
 
     # now create the expected tasks and the state is good again
-    comp_tasks = tasks(user=user, project=proj)
+    comp_tasks = await create_tasks(user=user, project=proj)
     response = await async_client.get(get_computation_url)
     assert response.status_code == status.HTTP_200_OK, response.text
     returned_computation = ComputationGet.model_validate(response.json())
@@ -990,8 +990,8 @@ async def test_get_computation_from_published_computation_task(
     registered_user: Callable[..., dict[str, Any]],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
-    runs: Callable[..., CompRunsAtDB],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
+    create_comp_run: Callable[..., Awaitable[CompRunsAtDB]],
     async_client: httpx.AsyncClient,
 ):
     user = registered_user()
@@ -1000,8 +1000,12 @@ async def test_get_computation_from_published_computation_task(
         project_id=proj.uuid,
         dag_adjacency_list=fake_workbench_adjacency,
     )
-    comp_tasks = tasks(user=user, project=proj, state=StateType.PUBLISHED, progress=0)
-    comp_runs = runs(user=user, project=proj, result=StateType.PUBLISHED)
+    comp_tasks = await create_tasks(
+        user=user, project=proj, state=StateType.PUBLISHED, progress=0
+    )
+    comp_runs = await create_comp_run(
+        user=user, project=proj, result=StateType.PUBLISHED
+    )
     assert comp_runs
     get_computation_url = httpx.URL(
         f"/v2/computations/{proj.uuid}?user_id={user['id']}"

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

@@ -9,7 +9,6 @@
 from uuid import uuid4
 
 import httpx
-from pydantic import TypeAdapter
 import pytest
 from faker import Faker
 from fastapi import FastAPI, status
@@ -22,6 +21,7 @@
 from models_library.projects import ProjectAtDB, ProjectID
 from models_library.projects_nodes_io import NodeID
 from models_library.users import UserID
+from pydantic import TypeAdapter
 from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
 from pytest_simcore.helpers.typing_env import EnvVarsDict
 from simcore_service_director_v2.core.settings import AppSettings
@@ -117,7 +117,7 @@ async def project_id(
     user: dict[str, Any],
     project: Callable[..., Awaitable[ProjectAtDB]],
     pipeline: Callable[..., CompPipelineAtDB],
-    tasks: Callable[..., list[CompTaskAtDB]],
+    create_tasks: Callable[..., Awaitable[list[CompTaskAtDB]]],
 ):
     """project uuid of a saved project (w/ tasks up-to-date)"""
 
@@ -130,7 +130,7 @@ async def project_id(
         dag_adjacency_list=fake_workbench_adjacency,
     )
     # insert tasks -> comp_tasks
-    comp_tasks = tasks(user=user, project=proj)
+    comp_tasks = await create_tasks(user=user, project=proj)
 
     return proj.uuid