ongoing tests

Author: sanderegg

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

@@ -35,6 +35,7 @@
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
 from models_library.rabbitmq_messages import (
+    ComputationalPipelineStatusMessage,
     InstrumentationRabbitMessage,
     RabbitResourceTrackingBaseMessage,
     RabbitResourceTrackingHeartbeatMessage,
@@ -198,6 +199,7 @@ async def _assert_publish_in_dask_backend(
     published_tasks: list[CompTaskAtDB],
     mocked_dask_client: mock.MagicMock,
     scheduler: BaseCompScheduler,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ) -> tuple[list[CompTaskAtDB], dict[NodeID, Callable[[], None]]]:
     expected_pending_tasks = [
         published_tasks[1],
@@ -285,6 +287,11 @@ async def _return_tasks_pending(job_ids: list[str]) -> list[RunningState]:
         where_statement=(comp_runs.c.user_id == published_project.project.prj_owner)
         & (comp_runs.c.project_uuid == f"{published_project.project.uuid}"),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     await assert_comp_tasks_and_comp_run_snapshot_tasks(
         sqlalchemy_async_engine,
         project_uuid=published_project.project.uuid,
@@ -336,6 +343,19 @@ async def resource_tracking_rabbit_client_parser(
     await client.unsubscribe(queue_name)
 
 
+@pytest.fixture
+async def computational_pipeline_rabbit_client_parser(
+    create_rabbitmq_client: Callable[[str], RabbitMQClient], mocker: MockerFixture
+) -> AsyncIterator[mock.AsyncMock]:
+    client = create_rabbitmq_client("computational_pipeline_pytest_consumer")
+    mock = mocker.AsyncMock(return_value=True)
+    queue_name, _ = await client.subscribe(
+        ComputationalPipelineStatusMessage.get_channel_name(), mock
+    )
+    yield mock
+    await client.unsubscribe(queue_name)
+
+
 async def _assert_message_received(
     mocked_message_parser: mock.AsyncMock,
     expected_call_count: int,
@@ -422,6 +442,7 @@ async def test_proper_pipeline_is_scheduled(  # noqa: PLR0915
     mocked_clean_task_output_and_log_files_if_invalid: mock.Mock,
     instrumentation_rabbit_client_parser: mock.AsyncMock,
     resource_tracking_rabbit_client_parser: mock.AsyncMock,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
     run_metadata: RunMetadataDict,
 ):
     with_disabled_auto_scheduling.assert_called_once()
@@ -438,6 +459,12 @@ async def test_proper_pipeline_is_scheduled(  # noqa: PLR0915
     )
     with_disabled_scheduler_publisher.assert_called()
 
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
+
     # -------------------------------------------------------------------------------
     # 1. first run will move comp_tasks to PENDING so the dask-worker can take them
     expected_pending_tasks, _ = await _assert_publish_in_dask_backend(
@@ -446,6 +473,7 @@ async def test_proper_pipeline_is_scheduled(  # noqa: PLR0915
         expected_published_tasks,
         mocked_dask_client,
         scheduler_api,
+        computational_pipeline_rabbit_client_parser,
     )
 
     # -------------------------------------------------------------------------------
@@ -478,6 +506,11 @@ async def _return_1st_task_running(job_ids: list[str]) -> list[RunningState]:
             comp_runs.c.project_uuid == f"{published_project.project.uuid}",
         ),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     await assert_comp_tasks_and_comp_run_snapshot_tasks(
         sqlalchemy_async_engine,
         project_uuid=published_project.project.uuid,
@@ -854,6 +887,11 @@ async def _return_3rd_task_success(job_ids: list[str]) -> list[RunningState]:
             comp_runs.c.project_uuid == f"{published_project.project.uuid}",
         ),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
 
     await assert_comp_tasks_and_comp_run_snapshot_tasks(
         sqlalchemy_async_engine,
@@ -903,6 +941,7 @@ async def with_started_project(
     scheduler_api: BaseCompScheduler,
     instrumentation_rabbit_client_parser: mock.AsyncMock,
     resource_tracking_rabbit_client_parser: mock.AsyncMock,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ) -> RunningProject:
     with_disabled_auto_scheduling.assert_called_once()
     published_project = await publish_project()
@@ -916,7 +955,11 @@ async def with_started_project(
         run_metadata=run_metadata,
     )
     with_disabled_scheduler_publisher.assert_called_once()
-
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     #
     # 2. This runs the scheduler until the project is started scheduled in the back-end
     #
@@ -929,6 +972,7 @@ async def with_started_project(
         expected_published_tasks,
         mocked_dask_client,
         scheduler_api,
+        computational_pipeline_rabbit_client_parser,
     )
 
     #
@@ -966,6 +1010,11 @@ async def _return_1st_task_running(job_ids: list[str]) -> list[RunningState]:
             comp_runs.c.project_uuid == f"{published_project.project.uuid}",
         ),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     await assert_comp_tasks_and_comp_run_snapshot_tasks(
         sqlalchemy_async_engine,
         project_uuid=published_project.project.uuid,
@@ -1140,6 +1189,7 @@ async def test_broken_pipeline_configuration_is_not_scheduled_and_aborted(
     fake_workbench_adjacency: dict[str, Any],
     sqlalchemy_async_engine: AsyncEngine,
     run_metadata: RunMetadataDict,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ):
     """A pipeline which comp_tasks are missing should not be scheduled.
     It shall be aborted and shown as such in the comp_runs db"""
@@ -1172,6 +1222,11 @@ async def test_broken_pipeline_configuration_is_not_scheduled_and_aborted(
             & (comp_runs.c.project_uuid == f"{sleepers_project.uuid}"),
         )
     )[0]
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
 
     #
     # Trigger scheduling manually. since the pipeline is broken, it shall be aborted
@@ -1188,6 +1243,11 @@ async def test_broken_pipeline_configuration_is_not_scheduled_and_aborted(
         where_statement=(comp_runs.c.user_id == user["id"])
         & (comp_runs.c.project_uuid == f"{sleepers_project.uuid}"),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
 
 
 async def test_task_progress_triggers(
@@ -1201,6 +1261,7 @@ async def test_task_progress_triggers(
     mocked_parse_output_data_fct: mock.Mock,
     mocked_clean_task_output_and_log_files_if_invalid: mock.Mock,
     run_metadata: RunMetadataDict,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ):
     _with_mock_send_computation_tasks(published_project.tasks, mocked_dask_client)
     _run_in_db, expected_published_tasks = await _assert_start_pipeline(
@@ -1209,6 +1270,11 @@ async def test_task_progress_triggers(
         published_project=published_project,
         run_metadata=run_metadata,
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     # -------------------------------------------------------------------------------
     # 1. first run will move comp_tasks to PENDING so the dask-worker can take them
     expected_pending_tasks, _ = await _assert_publish_in_dask_backend(
@@ -1217,6 +1283,7 @@ async def test_task_progress_triggers(
         expected_published_tasks,
         mocked_dask_client,
         scheduler_api,
+        computational_pipeline_rabbit_client_parser,
     )
 
     # send some progress
@@ -1272,6 +1339,7 @@ async def test_handling_of_disconnected_scheduler_dask(
     published_project: PublishedProject,
     backend_error: ComputationalSchedulerError,
     run_metadata: RunMetadataDict,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ):
     # this will create a non connected backend issue that will trigger re-connection
     mocked_dask_client_send_task = mocker.patch(
@@ -1289,6 +1357,11 @@ async def test_handling_of_disconnected_scheduler_dask(
         run_metadata=run_metadata,
         use_on_demand_clusters=False,
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
 
     # since there is no cluster, there is no dask-scheduler,
     # the tasks shall all still be in PUBLISHED state now
@@ -1337,6 +1410,11 @@ async def test_handling_of_disconnected_scheduler_dask(
         expected_progress=1,
         run_id=run_in_db.run_id,
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
     # then we have another scheduler run
     await scheduler_api.apply(
         user_id=run_in_db.user_id,
@@ -1452,6 +1530,7 @@ async def test_handling_scheduled_tasks_after_director_reboots(
     mocked_parse_output_data_fct: mock.Mock,
     mocked_clean_task_output_fct: mock.Mock,
     reboot_state: RebootState,
+    computational_pipeline_rabbit_client_parser: mock.AsyncMock,
 ):
     """After the dask client is rebooted, or that the director-v2 reboots the dv-2 internal scheduler
     shall continue scheduling correctly. Even though the task might have continued to run
@@ -1534,6 +1613,11 @@ async def mocked_get_task_result(_job_id: str) -> TaskOutputData:
             comp_runs.c.project_uuid == f"{running_project.project.uuid}",
         ),
     )
+    await _assert_message_received(
+        computational_pipeline_rabbit_client_parser,
+        1,
+        ComputationalPipelineStatusMessage.model_validate_json,
+    )
 
 
 async def test_handling_cancellation_of_jobs_after_reboot(