diff --git a/bec_widgets/utils/bec_dispatcher.py b/bec_widgets/utils/bec_dispatcher.py index 6c2f64e0..93dc50f5 100644 --- a/bec_widgets/utils/bec_dispatcher.py +++ b/bec_widgets/utils/bec_dispatcher.py @@ -259,50 +259,59 @@ class BECDispatcher: def disconnect_slot( self, slot: Callable, - topics: EndpointInfo | str | list[EndpointInfo] | list[str], + topics: EndpointInfo | str | list[EndpointInfo] | list[str] | None = None, cb_info: dict | None = None, ): """ - Disconnect a slot from a topic. + Disconnect a slot from topics. - Among the wrappers matching the callback (and ``cb_info``, when given), the one - actually subscribed to the requested topics is released. If none of them is - subscribed to these topics, nothing is released and a warning names the actual - subscriptions — a blind unregister on the wrong wrapper would silently no-op. + Among the wrappers matching the callback (and ``cb_info``, when given), every + wrapper subscribed to any of the requested topics is released from exactly those + topics — a topic list spanning several registrations of the same slot releases + all of them. Without ``topics``, the slot is disconnected from everything it is + subscribed to. If none of the wrappers is subscribed to the requested topics, + nothing is released and a warning names the actual subscriptions — a blind + unregister on the wrong wrapper would silently no-op. Args: slot(Callable): The slot to disconnect - topics EndpointInfo | str | list[EndpointInfo] | list[str]: A topic or list of topics to unsub from. + topics(EndpointInfo | str | list[EndpointInfo] | list[str] | None): A topic or + list of topics to unsub from. None disconnects the slot from all its topics. cb_info(dict | None): When the same slot was registered multiple times with different cb_info payloads (e.g. per-signal async subscriptions), pass the payload to narrow the candidate registrations before the topic match. """ - # find the right slot to disconnect from ; # slot callbacks are wrapped in QtThreadSafeCallback objects, # but the slot we receive here is the original callable - topics_str, _ = self.client.connector.extract_raw_endpoints_from_info(topics) - requested = set(topics_str) candidates = [ connected_slot - for connected_slot in self._registered_slots.values() + for connected_slot in list(self._registered_slots.values()) if connected_slot.cb == slot and (cb_info is None or connected_slot.cb_info == cb_info) ] - # Among wrappers for this callback, pick the one actually subscribed to the - # requested topics; a bare unregister on the wrong wrapper would silently no-op - # and leave the real subscription alive. - connected_slot = next((s for s in candidates if s.topics & requested), None) - if connected_slot is None: - if candidates: - logger.warning( - f"disconnect_slot({slot!r}): no registration matches topics " - f"{sorted(requested)}; the slot is subscribed to " - f"{sorted(set().union(*(s.topics for s in candidates)))}. Nothing released." - ) + if topics is None: + for connected_slot in candidates: + self._release_slot(connected_slot) return - self.client.connector.unregister(topics, cb=connected_slot) - self._registered_slots[connected_slot].topics.difference_update(requested) - if not self._registered_slots[connected_slot].topics: - del self._registered_slots[connected_slot] + topics_str, _ = self.client.connector.extract_raw_endpoints_from_info(topics) + requested = set(topics_str) + matched = False + for connected_slot in candidates: + overlap = connected_slot.topics & requested + if not overlap: + continue + matched = True + # unregister only what this wrapper actually holds; the same requested list + # may span several wrappers of the same callback + self.client.connector.unregister(list(overlap), cb=connected_slot) + connected_slot.topics.difference_update(overlap) + if not connected_slot.topics: + self._registered_slots.pop(connected_slot, None) + if not matched and candidates: + logger.warning( + f"disconnect_slot({slot!r}): no registration matches topics " + f"{sorted(requested)}; the slot is subscribed to " + f"{sorted(set().union(*(s.topics for s in candidates)))}. Nothing released." + ) def disconnect_topics(self, topics: str | list): """ diff --git a/tests/unit_tests/test_bec_dispatcher.py b/tests/unit_tests/test_bec_dispatcher.py index 0f6392a5..d172ddeb 100644 --- a/tests/unit_tests/test_bec_dispatcher.py +++ b/tests/unit_tests/test_bec_dispatcher.py @@ -320,3 +320,85 @@ def test_stop_cli_server_is_idempotent(bec_dispatcher): bec_dispatcher.stop_cli_server() bec_dispatcher.stop_cli_server() mock_logger.error.assert_not_called() + + +@pytest.mark.parametrize("topics_msg_list", [(("topic1", dummy_msg), ("topic2", dummy_msg))]) +def test_dispatcher_disconnect_topic_list_single_wrapper( + bec_dispatcher_w_connector, qtbot, send_msg_event +): + # One wrapper holding several topics: a list disconnect releases them all at once. + bec_dispatcher = bec_dispatcher_w_connector + cb1 = mock.Mock(spec=[]) + try: + bec_dispatcher.connect_slot(cb1, "topic1") + bec_dispatcher.connect_slot(cb1, "topic2") + assert len(bec_dispatcher.client.connector._managed_connection._topics_cb) == 2 + + bec_dispatcher.disconnect_slot(cb1, ["topic1", "topic2"]) + assert len(bec_dispatcher.client.connector._managed_connection._topics_cb) == 0 + finally: + # unblock the fixture's message generator so connector.shutdown() can join + send_msg_event.set() + qtbot.wait(10) + cb1.assert_not_called() + + +@pytest.mark.parametrize("topics_msg_list", [(("topic1", dummy_msg), ("topic2", dummy_msg))]) +def test_dispatcher_disconnect_topic_list_spans_wrappers( + bec_dispatcher_w_connector, qtbot, send_msg_event +): + # The same callback registered twice with different cb_info produces two wrappers, + # each holding one of the requested topics: a list disconnect must release BOTH, + # not only the first wrapper that overlaps. + bec_dispatcher = bec_dispatcher_w_connector + cb1 = mock.Mock(spec=[]) + try: + bec_dispatcher.connect_slot(cb1, "topic1", cb_info={"scan": "a"}) + bec_dispatcher.connect_slot(cb1, "topic2", cb_info={"scan": "b"}) + assert len(bec_dispatcher.client.connector._managed_connection._topics_cb) == 2 + + bec_dispatcher.disconnect_slot(cb1, ["topic1", "topic2"]) + assert len(bec_dispatcher.client.connector._managed_connection._topics_cb) == 0 + assert not any( + s.cb == cb1 for s in bec_dispatcher._registered_slots.values() + ), "all wrappers of the slot must be dropped" + finally: + send_msg_event.set() + qtbot.wait(10) + cb1.assert_not_called() + + +@pytest.mark.parametrize( + "topics_msg_list", [(("topic1", dummy_msg), ("topic2", dummy_msg), ("topic3", dummy_msg))] +) +def test_dispatcher_disconnect_without_topics_releases_slot_everywhere( + bec_dispatcher_w_connector, qtbot, send_msg_event +): + # Omitting topics disconnects the slot from everything it is subscribed to, + # across all its wrappers — while other slots stay untouched. + bec_dispatcher = bec_dispatcher_w_connector + cb1 = mock.Mock(spec=[]) + cb2 = mock.Mock(spec=[]) + try: + bec_dispatcher.connect_slot(cb1, "topic1", cb_info={"scan": "a"}) + bec_dispatcher.connect_slot(cb1, "topic2", cb_info={"scan": "a"}) + bec_dispatcher.connect_slot(cb1, "topic3", cb_info={"scan": "b"}) + bec_dispatcher.connect_slot(cb2, "topic1") + assert len(bec_dispatcher.client.connector._managed_connection._topics_cb) == 3 + + bec_dispatcher.disconnect_slot(cb1) + + remaining = bec_dispatcher.client.connector._managed_connection._topics_cb + assert list(remaining) == ["topic1"], "only cb2's topic1 subscription remains" + assert not any(s.cb == cb1 for s in bec_dispatcher._registered_slots.values()) + assert any(s.cb == cb2 for s in bec_dispatcher._registered_slots.values()) + finally: + send_msg_event.set() + qtbot.waitUntil(lambda: cb2.call_count == 1, timeout=2000) + cb1.assert_not_called() + + bec_dispatcher.disconnect_slot(cb2) + # the fixture's generator delivers messages for already-released topics, which + # auto-creates empty defaultdict keys — assert no callbacks remain instead + assert not any(bec_dispatcher.client.connector._managed_connection._topics_cb.values()) + assert not any(s.cb == cb2 for s in bec_dispatcher._registered_slots.values())