use begin

Author: sanderegg

services/director-v2/tests/unit/_helpers.py

@@ -18,6 +18,7 @@
 
 @dataclass
 class PublishedProject:
+    user: dict[str, Any]
     project: ProjectAtDB
     pipeline: CompPipelineAtDB
     tasks: list[CompTaskAtDB]

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

@@ -1,7 +1,19 @@
+# pylint: disable=no-value-for-parameter
+# pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
+# pylint: disable=too-many-arguments
+# pylint: disable=unused-argument
+# pylint: disable=unused-variable
+
+from typing import Awaitable, Callable
+
 import pytest
+from _helpers import PublishedProject
+from faker import Faker
 from models_library.projects import ProjectID
 from models_library.users import UserID
 from simcore_service_director_v2.core.errors import ComputationalRunNotFoundError
+from simcore_service_director_v2.models.comp_runs import CompRunsAtDB
 from simcore_service_director_v2.modules.db.repositories.comp_runs import (
     CompRunsRepository,
 )
@@ -14,9 +26,38 @@
 ]
 
 
-async def test_get(aiopg_engine, user_id: UserID, project_id: ProjectID):
+@pytest.fixture
+def fake_user_id(faker: Faker) -> UserID:
+    return faker.pyint(min_value=1)
+
+
+@pytest.fixture
+def fake_project_id(faker: Faker) -> ProjectID:
+    return ProjectID(f"{faker.uuid4(cast_to=None)}")
+
+
+async def test_get(
+    aiopg_engine,
+    fake_user_id: UserID,
+    fake_project_id: ProjectID,
+    publish_project: Callable[[], Awaitable[PublishedProject]],
+    create_comp_run: Callable[..., Awaitable[CompRunsAtDB]],
+):
+    with pytest.raises(ComputationalRunNotFoundError):
+        await CompRunsRepository(aiopg_engine).get(fake_user_id, fake_project_id)
+
+    published_project = await publish_project()
+    assert published_project.project.prj_owner
+    # there is still no comp run created
     with pytest.raises(ComputationalRunNotFoundError):
-        await CompRunsRepository(aiopg_engine).get(user_id, project_id)
+        await CompRunsRepository(aiopg_engine).get(
+            published_project.project.prj_owner, published_project.project.uuid
+        )
+
+    comp_run = await create_comp_run(published_project.user, published_project.project)
+    await CompRunsRepository(aiopg_engine).get(
+        published_project.project.prj_owner, published_project.project.uuid
+    )
 
 
 async def test_list():

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

@@ -1046,6 +1046,7 @@ async def _return_1st_task_running(job_ids: list[str]) -> list[DaskClientTaskSta
     assert messages[0].node_id == exp_started_task.node_id
 
     return RunningProject(
+        published_project.user,
         published_project.project,
         published_project.pipeline,
         tasks_in_db,

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

@@ -67,7 +67,7 @@ async def _(**pipeline_kwargs) -> CompPipelineAtDB:
     yield _
 
     # cleanup
-    async with sqlalchemy_async_engine.connect() as conn:
+    async with sqlalchemy_async_engine.begin() as conn:
         await conn.execute(
             comp_pipeline.delete().where(
                 comp_pipeline.c.project_id.in_(created_pipeline_ids)
@@ -123,7 +123,7 @@ async def _(
                 ),
                 "node_class": to_node_class(node_data.key),
                 "internal_id": internal_id + 1,
-                "submit": datetime.datetime.now(tz=datetime.UTC),
+                "submit": datetime.datetime.now(),
                 "job_id": generate_dask_job_id(
                     service_key=node_data.key,
                     service_version=node_data.version,
@@ -133,7 +133,7 @@ async def _(
                 ),
             }
             task_config.update(**overrides_kwargs)
-            async with sqlalchemy_async_engine.connect() as conn:
+            async with sqlalchemy_async_engine.begin() as conn:
                 result = await conn.execute(
                     comp_tasks.insert()
                     .values(**task_config)
@@ -147,7 +147,7 @@ async def _(
     yield _
 
     # cleanup
-    async with sqlalchemy_async_engine.connect() as conn:
+    async with sqlalchemy_async_engine.begin() as conn:
         await conn.execute(
             comp_tasks.delete().where(comp_tasks.c.task_id.in_(created_task_ids))
         )
@@ -197,14 +197,14 @@ async def _(
     ) -> CompRunsAtDB:
         run_config = {
             "project_uuid": f"{project.uuid}",
-            "user_id": f"{user['id']}",
+            "user_id": user["id"],
             "iteration": 1,
             "result": StateType.NOT_STARTED,
             "metadata": run_metadata,
             "use_on_demand_clusters": False,
         }
         run_config.update(**run_kwargs)
-        async with sqlalchemy_async_engine.connect() as conn:
+        async with sqlalchemy_async_engine.begin() as conn:
             result = await conn.execute(
                 comp_runs.insert()
                 .values(**jsonable_encoder(run_config))
@@ -217,7 +217,7 @@ async def _(
     yield _
 
     # cleanup
-    async with sqlalchemy_async_engine.connect() as conn:
+    async with sqlalchemy_async_engine.begin() as conn:
         await conn.execute(
             comp_runs.delete().where(comp_runs.c.run_id.in_(created_run_ids))
         )
@@ -242,7 +242,7 @@ async def _(user: dict[str, Any], **cluster_kwargs) -> Cluster:
         new_cluster = Cluster.model_validate(cluster_config)
         assert new_cluster
 
-        async with sqlalchemy_async_engine.connect() as conn:
+        async with sqlalchemy_async_engine.begin() as conn:
             # insert basic cluster
             created_cluster = (
                 await conn.execute(
@@ -298,7 +298,7 @@ async def _(user: dict[str, Any], **cluster_kwargs) -> Cluster:
     yield _
 
     # cleanup
-    async with sqlalchemy_async_engine.connect() as conn:
+    async with sqlalchemy_async_engine.begin() as conn:
         await conn.execute(
             clusters.delete().where(clusters.c.id.in_(created_cluster_ids))
         )
@@ -318,6 +318,7 @@ async def publish_project(
     async def _() -> PublishedProject:
         created_project = await project(user, workbench=fake_workbench_without_outputs)
         return PublishedProject(
+            user=user,
             project=created_project,
             pipeline=await create_pipeline(
                 project_id=f"{created_project.uuid}",
@@ -352,6 +353,7 @@ async def running_project(
     created_project = await project(user, workbench=fake_workbench_without_outputs)
     now_time = arrow.utcnow().datetime
     return RunningProject(
+        user=user,
         project=created_project,
         pipeline=await create_pipeline(
             project_id=f"{created_project.uuid}",
@@ -388,6 +390,7 @@ async def running_project_mark_for_cancellation(
     created_project = await project(user, workbench=fake_workbench_without_outputs)
     now_time = arrow.utcnow().datetime
     return RunningProject(
+        user=user,
         project=created_project,
         pipeline=await create_pipeline(
             project_id=f"{created_project.uuid}",