diff --git a/src/openai/_event_handler.py b/src/openai/_event_handler.py index bbcbd22a6b..c141bdcef3 100644 --- a/src/openai/_event_handler.py +++ b/src/openai/_event_handler.py @@ -42,7 +42,7 @@ def remove(self, event_type: str, handler: EventHandler) -> None: handlers = self._handlers.get(event_type) if handlers is not None: for index, registration in enumerate(handlers): - if registration.handler == handler: + if registration.handler is handler or registration.handler == handler: del handlers[index] break finally: diff --git a/tests/test_event_handler.py b/tests/test_event_handler.py new file mode 100644 index 0000000000..fc71e31f61 --- /dev/null +++ b/tests/test_event_handler.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from openai._event_handler import EventHandlerRegistry + + +def test_once_registration_is_scoped_to_its_event() -> None: + def handler(_event: object) -> None: + pass + + registry = EventHandlerRegistry() + registry.add("persistent", handler) + registry.add("one-shot", handler, once=True) + + assert registry.get_handlers("persistent") == [handler] + assert registry.get_handlers("persistent") == [handler] + assert registry.get_handlers("one-shot") == [handler] + assert registry.get_handlers("one-shot") == [] + + +def test_remove_only_removes_the_first_matching_registration() -> None: + def handler(_event: object) -> None: + pass + + registry = EventHandlerRegistry() + registry.add("event", handler) + registry.add("event", handler, once=True) + + registry.remove("event", handler) + + assert registry.get_handlers("event") == [handler] + assert registry.get_handlers("event") == [] + + +def test_remove_prefers_identity_to_overloaded_equality() -> None: + class Handler: + def __call__(self, _event: object) -> None: + pass + + def __eq__(self, other: object) -> bool: + del other + raise AssertionError("identity matches must not invoke __eq__") + + handler = Handler() + registry = EventHandlerRegistry() + registry.add("event", handler) + + registry.remove("event", handler) + + assert registry.has_handlers("event") is False + + +def test_merge_preserves_each_registration_mode() -> None: + def handler(_event: object) -> None: + pass + + source = EventHandlerRegistry() + target = EventHandlerRegistry() + source.add("event", handler) + source.add("event", handler, once=True) + + source.merge_into(target) + + assert source.has_handlers("event") is False + assert target.get_handlers("event") == [handler, handler] + assert target.get_handlers("event") == [handler]