more reliable

Author: sanderegg

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

@@ -9,6 +9,7 @@
 
 
 import asyncio
+import contextlib
 import datetime
 import logging
 from collections.abc import AsyncIterator, Awaitable, Callable
@@ -42,6 +43,12 @@
     CompRunsRepository,
 )
 from sqlalchemy.ext.asyncio import AsyncEngine
+from tenacity import (
+    RetryError,
+    retry,
+    retry_if_exception_type,
+    retry_unless_exception_type,
+)
 
 pytest_simcore_core_services_selection = ["postgres", "rabbit", "redis"]
 pytest_simcore_ops_services_selection = ["adminer", "redis-commander"]
@@ -102,7 +109,7 @@ async def test_schedule_all_pipelines_empty_db(
     await schedule_all_pipelines(initialized_app)
 
     # check nothing was distributed
-    scheduler_rabbit_client_parser.assert_not_called()
+    _assert_scheduler_client_not_called(scheduler_rabbit_client_parser)
 
     # check comp_runs is still empty
     await assert_comp_runs_empty(sqlalchemy_async_engine)
@@ -142,6 +149,35 @@ async def slow_limited_gather(*args, **kwargs):
     mock_function.assert_called_once()
 
 
+@retry(
+    retry=retry_if_exception_type(AssertionError),
+    stop=retry.stop_after_attempt(3),
+    wait=retry.wait_fixed(0.5),
+    reraise=True,
+)
+def _assert_scheduler_client_called_once_with(
+    scheduler_rabbit_client_parser: mock.AsyncMock,
+    expected_message: SchedulePipelineRabbitMessage,
+):
+    scheduler_rabbit_client_parser.assert_called_once_with(expected_message.body())
+
+
+def _assert_scheduler_client_not_called(
+    scheduler_rabbit_client_parser: mock.AsyncMock,
+):
+    @retry(
+        retry=retry_unless_exception_type(AssertionError),
+        stop=retry.stop_after_attempt(2),
+        wait=retry.wait_fixed(1),
+        reraise=True,
+    )  # pylint: disable=unused-variable
+    def _():
+        scheduler_rabbit_client_parser.assert_not_called()
+
+    with contextlib.suppress(RetryError):
+        _()
+
+
 async def test_schedule_all_pipelines(
     with_disabled_auto_scheduling: mock.Mock,
     with_disabled_scheduler_worker: mock.Mock,
@@ -164,12 +200,13 @@ async def test_schedule_all_pipelines(
         collection_run_id=fake_collection_run_id,
     )
     # this directly schedule a new pipeline
-    scheduler_rabbit_client_parser.assert_called_once_with(
+    _assert_scheduler_client_called_once_with(
+        scheduler_rabbit_client_parser,
         SchedulePipelineRabbitMessage(
             user_id=published_project.project.prj_owner,
             project_id=published_project.project.uuid,
             iteration=1,
-        ).body()
+        ),
     )
     scheduler_rabbit_client_parser.reset_mock()
     comp_run = (await assert_comp_runs(sqlalchemy_async_engine, expected_total=1))[0]
@@ -186,7 +223,7 @@ async def test_schedule_all_pipelines(
 
     # this will now not schedule the pipeline since it was already scheduled
     await schedule_all_pipelines(initialized_app)
-    scheduler_rabbit_client_parser.assert_not_called()
+    _assert_scheduler_client_not_called(scheduler_rabbit_client_parser)
     comp_runs = await assert_comp_runs(sqlalchemy_async_engine, expected_total=1)
     comp_run = comp_runs[0]
     assert comp_run.scheduled
@@ -206,12 +243,13 @@ async def test_schedule_all_pipelines(
     # now we schedule a pipeline again, but we wait for the scheduler interval to pass
     # this will trigger a new schedule
     await schedule_all_pipelines(initialized_app)
-    scheduler_rabbit_client_parser.assert_called_once_with(
+    _assert_scheduler_client_called_once_with(
+        scheduler_rabbit_client_parser,
         SchedulePipelineRabbitMessage(
             user_id=published_project.project.prj_owner,
             project_id=published_project.project.uuid,
             iteration=1,
-        ).body()
+        ),
     )
     scheduler_rabbit_client_parser.reset_mock()
     comp_runs = await assert_comp_runs(sqlalchemy_async_engine, expected_total=1)
@@ -229,12 +267,13 @@ async def test_schedule_all_pipelines(
         project_id=published_project.project.uuid,
     )
     await schedule_all_pipelines(initialized_app)
-    scheduler_rabbit_client_parser.assert_called_once_with(
+    _assert_scheduler_client_called_once_with(
+        scheduler_rabbit_client_parser,
         SchedulePipelineRabbitMessage(
             user_id=published_project.project.prj_owner,
             project_id=published_project.project.uuid,
             iteration=1,
-        ).body()
+        ),
     )
     scheduler_rabbit_client_parser.reset_mock()
     comp_runs = await assert_comp_runs(sqlalchemy_async_engine, expected_total=1)
@@ -267,12 +306,13 @@ async def test_schedule_all_pipelines_logs_error_if_it_find_old_pipelines(
         collection_run_id=fake_collection_run_id,
     )
     # this directly schedule a new pipeline
-    scheduler_rabbit_client_parser.assert_called_once_with(
+    _assert_scheduler_client_called_once_with(
+        scheduler_rabbit_client_parser,
         SchedulePipelineRabbitMessage(
             user_id=published_project.project.prj_owner,
             project_id=published_project.project.uuid,
             iteration=1,
-        ).body()
+        ),
     )
     scheduler_rabbit_client_parser.reset_mock()
     comp_run = (await assert_comp_runs(sqlalchemy_async_engine, expected_total=1))[0]
@@ -288,7 +328,7 @@ async def test_schedule_all_pipelines_logs_error_if_it_find_old_pipelines(
 
     # this will now not schedule the pipeline since it was already scheduled
     await schedule_all_pipelines(initialized_app)
-    scheduler_rabbit_client_parser.assert_not_called()
+    _assert_scheduler_client_not_called(scheduler_rabbit_client_parser)
     comp_runs = await assert_comp_runs(sqlalchemy_async_engine, expected_total=1)
     comp_run = comp_runs[0]
     assert comp_run.scheduled == start_schedule_time, "scheduled time changed!"
@@ -308,12 +348,13 @@ async def test_schedule_all_pipelines_logs_error_if_it_find_old_pipelines(
         assert (
             "found 1 lost pipelines, they will be re-scheduled now" in caplog.messages
         )
-    scheduler_rabbit_client_parser.assert_called_once_with(
+    _assert_scheduler_client_called_once_with(
+        scheduler_rabbit_client_parser,
         SchedulePipelineRabbitMessage(
             user_id=published_project.project.prj_owner,
             project_id=published_project.project.uuid,
             iteration=1,
-        ).body()
+        ),
     )
     scheduler_rabbit_client_parser.reset_mock()
     comp_runs = await assert_comp_runs(sqlalchemy_async_engine, expected_total=1)
@@ -353,7 +394,7 @@ async def test_empty_pipeline_is_not_scheduled(
             collection_run_id=fake_collection_run_id,
         )
     await assert_comp_runs_empty(sqlalchemy_async_engine)
-    scheduler_rabbit_client_parser.assert_not_called()
+    _assert_scheduler_client_not_called(scheduler_rabbit_client_parser)
 
     # create the empty pipeline now
     await create_pipeline(project_id=f"{empty_project.uuid}")
@@ -371,4 +412,4 @@ async def test_empty_pipeline_is_not_scheduled(
     assert len(caplog.records) == 1
     assert "no computational dag defined" in caplog.records[0].message
     await assert_comp_runs_empty(sqlalchemy_async_engine)
-    scheduler_rabbit_client_parser.assert_not_called()
+    _assert_scheduler_client_not_called(scheduler_rabbit_client_parser)