More test fixups.

api/tests/opentrons/protocol_runner/test_legacy_context_plugin.py

@@ -3,7 +3,7 @@
 from anyio import to_thread
 from decoy import Decoy, matchers
 from datetime import datetime
-from typing import Callable
+from typing import Any, Callable, NamedTuple, TypeVar
 
 from opentrons.commands.types import CommandMessage as LegacyCommand, PauseMessage
 from opentrons.hardware_control import API as HardwareAPI
@@ -79,6 +79,72 @@ def subject(
     return plugin
 
 
+_UnsubscribeCallback = Callable[[], None]
+_PayloadT = TypeVar("_PayloadT")
+_Handler = Callable[[_PayloadT], None]
+
+
+class _UnsubscribeCallbackSet(NamedTuple):
+    command_unsubscribe: _UnsubscribeCallback
+    labware_unsubscribe: _UnsubscribeCallback
+    instrument_unsubscribe: _UnsubscribeCallback
+    module_unsubscribe: _UnsubscribeCallback
+
+
+class _HandlerCaptorSet(NamedTuple):
+    command_handler_captor: Any
+    labware_handler_captor: Any
+    instrument_handler_captor: Any
+    module_handler_captor: Any
+
+
+class _Result(NamedTuple):
+    handler_captors: _HandlerCaptorSet
+    unsubscribe_callbacks: _UnsubscribeCallbackSet
+
+
+def _mock_subscribes(decoy: Decoy, legacy_context: LegacyProtocolContext) -> _Result:
+    handler_captors = _HandlerCaptorSet(
+        command_handler_captor=matchers.Captor(),
+        labware_handler_captor=matchers.Captor(),
+        instrument_handler_captor=matchers.Captor(),
+        module_handler_captor=matchers.Captor(),
+    )
+
+    unsubscribe_callbacks = _UnsubscribeCallbackSet(
+        command_unsubscribe=decoy.mock(),
+        labware_unsubscribe=decoy.mock(),
+        instrument_unsubscribe=decoy.mock(),
+        module_unsubscribe=decoy.mock(),
+    )
+
+    decoy.when(
+        legacy_context.broker.subscribe(
+            topic="command", handler=handler_captors.command_handler_captor
+        )
+    ).then_return(unsubscribe_callbacks.command_unsubscribe)
+    decoy.when(
+        legacy_context.labware_load_broker.subscribe(
+            callback=handler_captors.labware_handler_captor
+        )
+    ).then_return(unsubscribe_callbacks.labware_unsubscribe)
+    decoy.when(
+        legacy_context.instrument_load_broker.subscribe(
+            callback=handler_captors.instrument_handler_captor
+        )
+    ).then_return(unsubscribe_callbacks.instrument_unsubscribe)
+    decoy.when(
+        legacy_context.module_load_broker.subscribe(
+            callback=handler_captors.module_handler_captor
+        )
+    ).then_return(unsubscribe_callbacks.module_unsubscribe)
+
+    return _Result(
+        handler_captors=handler_captors,
+        unsubscribe_callbacks=unsubscribe_callbacks,
+    )
+
+
 def test_play_action(
     decoy: Decoy,
     hardware_api: HardwareAPI,
@@ -151,14 +217,13 @@ async def test_main_broker_messages(
     subject: LegacyContextPlugin,
 ) -> None:
     """It should dispatch commands from main broker messages."""
-    await subject.setup()
+    handler_captors, _ = _mock_subscribes(decoy=decoy, legacy_context=legacy_context)
 
-    handler_captor = matchers.Captor()
-    decoy.verify(
-        legacy_context.broker.subscribe(topic="command", handler=handler_captor)
-    )
+    await subject.setup()
 
-    handler: Callable[[LegacyCommand], None] = handler_captor.value
+    handler: Callable[
+        [LegacyCommand], None
+    ] = handler_captors.command_handler_captor.value
 
     legacy_command: PauseMessage = {
         "$": "before",
@@ -197,13 +262,13 @@ async def test_labware_load_broker_messages(
     minimal_labware_def: LabwareDefinitionDict,
 ) -> None:
     """It should dispatch commands from labware load broker messages."""
-    subject.setup()
+    handler_captors, _ = _mock_subscribes(decoy=decoy, legacy_context=legacy_context)
 
-    handler_captor = matchers.Captor()
-
-    decoy.verify(legacy_context.labware_load_broker.subscribe(callback=handler_captor))
+    await subject.setup()
 
-    handler: Callable[[LegacyLabwareLoadInfo], None] = handler_captor.value
+    handler: Callable[
+        [LegacyLabwareLoadInfo], None
+    ] = handler_captors.labware_handler_captor.value
 
     labware_load_info = LegacyLabwareLoadInfo(
         labware_definition=minimal_labware_def,
@@ -229,6 +294,8 @@ async def test_labware_load_broker_messages(
 
     await to_thread.run_sync(handler, labware_load_info)
 
+    await subject.teardown()
+
     decoy.verify(
         action_dispatcher.dispatch(pe_actions.UpdateCommandAction(engine_command))
     )
@@ -242,15 +309,13 @@ async def test_instrument_load_broker_messages(
     subject: LegacyContextPlugin,
 ) -> None:
     """It should dispatch commands from instrument load broker messages."""
-    subject.setup()
-
-    handler_captor = matchers.Captor()
+    handler_captors, _ = _mock_subscribes(decoy=decoy, legacy_context=legacy_context)
 
-    decoy.verify(
-        legacy_context.instrument_load_broker.subscribe(callback=handler_captor)
-    )
+    await subject.setup()
 
-    handler: Callable[[LegacyInstrumentLoadInfo], None] = handler_captor.value
+    handler: Callable[
+        [LegacyInstrumentLoadInfo], None
+    ] = handler_captors.instrument_handler_captor.value
 
     instrument_load_info = LegacyInstrumentLoadInfo(
         instrument_load_name="some_load_name", mount=Mount.LEFT
@@ -272,6 +337,8 @@ async def test_instrument_load_broker_messages(
 
     await to_thread.run_sync(handler, instrument_load_info)
 
+    await subject.teardown()
+
     decoy.verify(
         action_dispatcher.dispatch(pe_actions.UpdateCommandAction(engine_command))
     )
@@ -285,13 +352,13 @@ async def test_module_load_broker_messages(
     subject: LegacyContextPlugin,
 ) -> None:
     """It should dispatch commands from module load broker messages."""
-    await subject.setup()
+    handler_captors, _ = _mock_subscribes(decoy=decoy, legacy_context=legacy_context)
 
-    handler_captor = matchers.Captor()
-
-    decoy.verify(legacy_context.module_load_broker.subscribe(callback=handler_captor))
+    await subject.setup()
 
-    handler: Callable[[LegacyModuleLoadInfo], None] = handler_captor.value
+    handler: Callable[
+        [LegacyModuleLoadInfo], None
+    ] = handler_captors.module_handler_captor.value
 
     module_load_info = LegacyModuleLoadInfo(
         module_model=LegacyMagneticModuleModel.MAGNETIC_V2,
@@ -313,8 +380,8 @@ async def test_module_load_broker_messages(
 
     await to_thread.run_sync(handler, module_load_info)
 
+    await subject.teardown()
+
     decoy.verify(
         action_dispatcher.dispatch(pe_actions.UpdateCommandAction(engine_command))
     )
-
-    # Forgetting to do teardown here. What happens?

api/tests/opentrons/protocol_runner/test_protocol_runner.py

@@ -201,7 +201,7 @@ async def test_run(
 
 
 @pytest.mark.xfail(raises=NotImplementedError, strict=True)
-def test_load_json(
+async def test_load_json(
     decoy: Decoy,
     json_file_reader: JsonFileReader,
     json_command_translator: JsonCommandTranslator,
@@ -229,7 +229,7 @@ def test_load_json(
     decoy.when(json_file_reader.read(json_protocol_source)).then_return(json_protocol)
     decoy.when(json_command_translator.translate(json_protocol)).then_return(commands)
 
-    subject.load(json_protocol_source)
+    await subject.load(json_protocol_source)
 
     decoy.verify(
         protocol_engine.add_command(
@@ -246,7 +246,7 @@ def test_load_json(
     )
 
 
-def test_load_python(
+async def test_load_python(
     decoy: Decoy,
     python_file_reader: PythonFileReader,
     python_context_creator: PythonContextCreator,
@@ -275,7 +275,7 @@ def test_load_python(
         protocol_context
     )
 
-    subject.load(python_protocol_source)
+    await subject.load(python_protocol_source)
 
     decoy.verify(
         task_queue.set_run_func(
@@ -334,7 +334,7 @@ async def test_load_legacy_python(
 
     decoy.verify(
         protocol_engine.add_labware_definition(labware_definition),
-        protocol_engine.add_plugin(matchers.IsA(LegacyContextPlugin)),
+        await protocol_engine.add_plugin(matchers.IsA(LegacyContextPlugin)),
         task_queue.set_run_func(
             func=legacy_executor.execute,
             protocol=legacy_protocol,
@@ -386,7 +386,7 @@ async def test_load_legacy_json(
 
     decoy.verify(
         protocol_engine.add_labware_definition(labware_definition),
-        protocol_engine.add_plugin(matchers.IsA(LegacyContextPlugin)),
+        await protocol_engine.add_plugin(matchers.IsA(LegacyContextPlugin)),
         task_queue.set_run_func(
             func=legacy_executor.execute,
             protocol=legacy_protocol,