Fix system and client tests

tests/system_tests/test_blueapi_system.py

@@ -123,6 +123,11 @@ def clean_existing_tasks(client: BlueapiClient):
     yield
 
 
+@pytest.fixture
+def instrument_session() -> str:
+    return "cm12345-1"
+
+
 @pytest.mark.xfail(reason=_REQUIRES_AUTH_MESSAGE)
 def test_cannot_access_endpoints(
     client_without_auth: BlueapiClient, blueapi_client_get_methods: list[str]
@@ -229,24 +234,26 @@ def test_delete_non_existent_task(client: BlueapiClient):
         client.clear_task("Not-exists")
 
 
-def test_put_worker_task(client: BlueapiClient):
+def test_put_worker_task(client: BlueapiClient, instrument_session: str):
     created_task = client.create_task(_SIMPLE_TASK)
-    client.start_task(WorkerTask(task_id=created_task.task_id))
+    client.start_task(WorkerTask(task_id=created_task.task_id), instrument_session)
     active_task = client.get_active_task()
     assert active_task.task_id == created_task.task_id
     client.clear_task(created_task.task_id)
 
 
-def test_put_worker_task_fails_if_not_idle(client: BlueapiClient):
+def test_put_worker_task_fails_if_not_idle(
+    client: BlueapiClient, instrument_session: str
+):
     small_task = client.create_task(_SIMPLE_TASK)
     long_task = client.create_task(_LONG_TASK)
 
-    client.start_task(WorkerTask(task_id=long_task.task_id))
+    client.start_task(WorkerTask(task_id=long_task.task_id), instrument_session)
     active_task = client.get_active_task()
     assert active_task.task_id == long_task.task_id
 
     with pytest.raises(BlueskyRemoteControlError) as exception:
-        client.start_task(WorkerTask(task_id=small_task.task_id))
+        client.start_task(WorkerTask(task_id=small_task.task_id), instrument_session)
     assert "<Response [409]>" in str(exception)
     client.abort()
     client.clear_task(small_task.task_id)
@@ -266,7 +273,7 @@ def test_set_state_transition_error(client: BlueapiClient):
     assert "<Response [400]>" in str(exception)
 
 
-def test_get_task_by_status(client: BlueapiClient):
+def test_get_task_by_status(client: BlueapiClient, instrument_session: str):
     task_1 = client.create_task(_SIMPLE_TASK)
     task_2 = client.create_task(_SIMPLE_TASK)
     task_by_pending = client.get_all_tasks()
@@ -278,10 +285,10 @@ def test_get_task_by_status(client: BlueapiClient):
         trackable_task = TypeAdapter(TrackableTask).validate_python(task)
         assert trackable_task.is_complete is False and trackable_task.is_pending is True
 
-    client.start_task(WorkerTask(task_id=task_1.task_id))
+    client.start_task(WorkerTask(task_id=task_1.task_id), instrument_session)
     while not client.get_task(task_1.task_id).is_complete:
         time.sleep(0.1)
-    client.start_task(WorkerTask(task_id=task_2.task_id))
+    client.start_task(WorkerTask(task_id=task_2.task_id), instrument_session)
     while not client.get_task(task_2.task_id).is_complete:
         time.sleep(0.1)
     task_by_completed = client.get_all_tasks()
@@ -297,13 +304,13 @@ def test_get_task_by_status(client: BlueapiClient):
     client.clear_task(task_id=task_2.task_id)
 
 
-def test_progress_with_stomp(client_with_stomp: BlueapiClient):
+def test_progress_with_stomp(client_with_stomp: BlueapiClient, instrument_session: str):
     all_events: list[AnyEvent] = []
 
     def on_event(event: AnyEvent):
         all_events.append(event)
 
-    client_with_stomp.run_task(_SIMPLE_TASK, on_event=on_event)
+    client_with_stomp.run_task(_SIMPLE_TASK, instrument_session, on_event=on_event)
     assert isinstance(all_events[0], WorkerEvent) and all_events[0].task_status
     task_id = all_events[0].task_status.task_id
     assert all_events == [
@@ -372,7 +379,9 @@ def test_delete_current_environment(client: BlueapiClient):
         ),
     ],
 )
-def test_plan_runs(client_with_stomp: BlueapiClient, task: Task):
-    final_event = client_with_stomp.run_task(task)
+def test_plan_runs(
+    client_with_stomp: BlueapiClient, task: Task, instrument_session: str
+):
+    final_event = client_with_stomp.run_task(task, instrument_session)
     assert final_event.is_complete() and not final_event.is_error()
     assert final_event.state is WorkerState.IDLE

tests/unit_tests/client/test_client.py

@@ -103,6 +103,11 @@ def client_with_events(mock_rest: Mock, mock_events: MagicMock):
     return BlueapiClient(rest=mock_rest, events=mock_events)
 
 
+@pytest.fixture
+def instrument_session() -> str:
+    return "cm12345-1"
+
+
 def test_get_plans(client: BlueapiClient):
     assert client.get_plans() == PLANS
 
@@ -194,61 +199,65 @@ def test_get_active_task(client: BlueapiClient):
     assert client.get_active_task() == ACTIVE_TASK
 
 
-def test_start_task(
-    client: BlueapiClient,
-    mock_rest: Mock,
-):
-    client.start_task(task=WorkerTask(task_id="bar"))
-    mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="bar"))
+def test_start_task(client: BlueapiClient, mock_rest: Mock, instrument_session: str):
+    client.start_task(WorkerTask(task_id="bar"), instrument_session)
+    mock_rest.update_worker_task.assert_called_once_with(
+        WorkerTask(task_id="bar"), instrument_session
+    )
 
 
 def test_start_nonexistent_task(
-    client: BlueapiClient,
-    mock_rest: Mock,
+    client: BlueapiClient, mock_rest: Mock, instrument_session: str
 ):
     mock_rest.update_worker_task.side_effect = KeyError("Not found")
     with pytest.raises(KeyError):
-        client.start_task(task=WorkerTask(task_id="bar"))
+        client.start_task(WorkerTask(task_id="bar"), instrument_session)
 
 
 def test_create_and_start_task_calls_both_creating_and_starting_endpoints(
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="baz")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="baz")
-    client.create_and_start_task(Task(name="baz"))
+    client.create_and_start_task(Task(name="baz"), instrument_session)
     mock_rest.create_task.assert_called_once_with(Task(name="baz"))
-    mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="baz"))
+    mock_rest.update_worker_task.assert_called_once_with(
+        WorkerTask(task_id="baz"), instrument_session
+    )
 
 
 def test_create_and_start_task_fails_if_task_creation_fails(
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     mock_rest.create_task.side_effect = BlueskyRemoteControlError("No can do")
     with pytest.raises(BlueskyRemoteControlError):
-        client.create_and_start_task(Task(name="baz"))
+        client.create_and_start_task(Task(name="baz"), instrument_session)
 
 
 def test_create_and_start_task_fails_if_task_id_is_wrong(
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="baz")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="bar")
     with pytest.raises(BlueskyRemoteControlError):
-        client.create_and_start_task(Task(name="baz"))
+        client.create_and_start_task(Task(name="baz"), instrument_session)
 
 
 def test_create_and_start_task_fails_if_task_start_fails(
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="baz")
     mock_rest.update_worker_task.side_effect = BlueskyRemoteControlError("No can do")
     with pytest.raises(BlueskyRemoteControlError):
-        client.create_and_start_task(Task(name="baz"))
+        client.create_and_start_task(Task(name="baz"), instrument_session)
 
 
 def test_get_environment(client: BlueapiClient):
@@ -378,34 +387,40 @@ def test_resume(
     )
 
 
-def test_cannot_run_task_without_message_bus(client: BlueapiClient):
+def test_cannot_run_task_without_message_bus(
+    client: BlueapiClient, instrument_session: str
+):
     with pytest.raises(
         RuntimeError,
         match="Cannot run plans without Stomp configuration to track progress",
     ):
-        client.run_task(Task(name="foo"))
+        client.run_task(Task(name="foo"), instrument_session)
 
 
 def test_run_task_sets_up_control(
     client_with_events: BlueapiClient,
     mock_rest: Mock,
     mock_events: MagicMock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="foo")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")
     ctx = Mock()
     ctx.correlation_id = "foo"
     mock_events.subscribe_to_all_events = lambda on_event: on_event(COMPLETE_EVENT, ctx)
 
-    client_with_events.run_task(Task(name="foo"))
+    client_with_events.run_task(Task(name="foo"), instrument_session)
     mock_rest.create_task.assert_called_once_with(Task(name="foo"))
-    mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="foo"))
+    mock_rest.update_worker_task.assert_called_once_with(
+        WorkerTask(task_id="foo"), instrument_session
+    )
 
 
 def test_run_task_fails_on_failing_event(
     client_with_events: BlueapiClient,
     mock_rest: Mock,
     mock_events: MagicMock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="foo")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")
@@ -416,7 +431,9 @@ def test_run_task_fails_on_failing_event(
 
     on_event = Mock()
     with pytest.raises(BlueskyStreamingError):
-        client_with_events.run_task(Task(name="foo"), on_event=on_event)
+        client_with_events.run_task(
+            Task(name="foo"), instrument_session, on_event=on_event
+        )
 
     on_event.assert_called_with(FAILED_EVENT)
 
@@ -441,6 +458,7 @@ def test_run_task_calls_event_callback(
     mock_rest: Mock,
     mock_events: MagicMock,
     test_event: AnyEvent,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="foo")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")
@@ -455,7 +473,9 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]):
     mock_events.subscribe_to_all_events = callback  # type: ignore
 
     mock_on_event = Mock()
-    client_with_events.run_task(Task(name="foo"), on_event=mock_on_event)
+    client_with_events.run_task(
+        Task(name="foo"), instrument_session, on_event=mock_on_event
+    )
 
     assert mock_on_event.mock_calls == [call(test_event), call(COMPLETE_EVENT)]
 
@@ -480,6 +500,7 @@ def test_run_task_ignores_non_matching_events(
     mock_rest: Mock,
     mock_events: MagicMock,
     test_event: AnyEvent,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="foo")  # type: ignore
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")  # type: ignore
@@ -494,7 +515,9 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]):
     mock_events.subscribe_to_all_events = callback
 
     mock_on_event = Mock()
-    client_with_events.run_task(Task(name="foo"), on_event=mock_on_event)
+    client_with_events.run_task(
+        Task(name="foo"), instrument_session, on_event=mock_on_event
+    )
 
     mock_on_event.assert_called_once_with(COMPLETE_EVENT)
 
@@ -566,20 +589,22 @@ def test_start_task_span_ok(
     exporter: JsonObjectSpanExporter,
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     with asserting_span_exporter(exporter, "start_task", "task"):
-        client.start_task(task=WorkerTask(task_id="bar"))
+        client.start_task(WorkerTask(task_id="bar"), instrument_session)
 
 
 def test_create_and_start_task_span_ok(
     exporter: JsonObjectSpanExporter,
     client: BlueapiClient,
     mock_rest: Mock,
+    instrument_session: str,
 ):
     mock_rest.create_task.return_value = TaskResponse(task_id="baz")
     mock_rest.update_worker_task.return_value = TaskResponse(task_id="baz")
     with asserting_span_exporter(exporter, "create_and_start_task", "task"):
-        client.create_and_start_task(Task(name="baz"))
+        client.create_and_start_task(Task(name="baz"), instrument_session)
 
 
 def test_get_environment_span_ok(
@@ -636,11 +661,11 @@ def test_resume_span_ok(
 
 
 def test_cannot_run_task_span_ok(
-    exporter: JsonObjectSpanExporter, client: BlueapiClient
+    exporter: JsonObjectSpanExporter, client: BlueapiClient, instrument_session: str
 ):
     with pytest.raises(
         RuntimeError,
         match="Cannot run plans without Stomp configuration to track progress",
     ):
         with asserting_span_exporter(exporter, "run_task"):
-            client.run_task(Task(name="foo"))
+            client.run_task(Task(name="foo"), instrument_session)

tests/unit_tests/test_cli.py

@@ -170,7 +170,8 @@ def test_submit_plan(runner: CliRunner):
 
     config_path = "tests/unit_tests/example_yaml/rest_config.yaml"
     runner.invoke(
-        main, ["-c", config_path, "controller", "run", "sleep", '{"time": 5}']
+        main,
+        ["-c", config_path, "controller", "run", "sleep", "cm12345-1", '{"time": 5}'],
     )
 
     assert response.call_count == 1