test: use scoped_handlers while listener Memory events

2026-01-08 23:58:34 +00:00 · 2025-07-01 23:07:56 -03:00
parent 15c5654a70
commit 174215980b
4 changed files with 119 additions and 124 deletions
--- a/tests/crew_test.py
+++ b/tests/crew_test.py
@@ -2499,58 +2499,59 @@ def test_using_contextual_memory():
 def test_memory_events_are_emitted():
    events = defaultdict(list)

-    @crewai_event_bus.on(MemorySaveStartedEvent)
-    def handle_memory_save_started(source, event):
-        events["MemorySaveStartedEvent"].append(event)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemorySaveStartedEvent)
+        def handle_memory_save_started(source, event):
+            events["MemorySaveStartedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveCompletedEvent)
-    def handle_memory_save_completed(source, event):
-        events["MemorySaveCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemorySaveCompletedEvent)
+        def handle_memory_save_completed(source, event):
+            events["MemorySaveCompletedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveFailedEvent)
-    def handle_memory_save_failed(source, event):
-        events["MemorySaveFailedEvent"].append(event)
+        @crewai_event_bus.on(MemorySaveFailedEvent)
+        def handle_memory_save_failed(source, event):
+            events["MemorySaveFailedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryStartedEvent)
-    def handle_memory_query_started(source, event):
-        events["MemoryQueryStartedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryStartedEvent)
+        def handle_memory_query_started(source, event):
+            events["MemoryQueryStartedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryCompletedEvent)
-    def handle_memory_query_completed(source, event):
-        events["MemoryQueryCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryCompletedEvent)
+        def handle_memory_query_completed(source, event):
+            events["MemoryQueryCompletedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryFailedEvent)
-    def handle_memory_query_failed(source, event):
-        events["MemoryQueryFailedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryFailedEvent)
+        def handle_memory_query_failed(source, event):
+            events["MemoryQueryFailedEvent"].append(event)

-    @crewai_event_bus.on(MemoryRetrievalStartedEvent)
-    def handle_memory_retrieval_started(source, event):
-        events["MemoryRetrievalStartedEvent"].append(event)
+        @crewai_event_bus.on(MemoryRetrievalStartedEvent)
+        def handle_memory_retrieval_started(source, event):
+            events["MemoryRetrievalStartedEvent"].append(event)

-    @crewai_event_bus.on(MemoryRetrievalCompletedEvent)
-    def handle_memory_retrieval_completed(source, event):
-        events["MemoryRetrievalCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemoryRetrievalCompletedEvent)
+        def handle_memory_retrieval_completed(source, event):
+            events["MemoryRetrievalCompletedEvent"].append(event)

-    math_researcher = Agent(
-        role="Researcher",
-        goal="You research about math.",
-        backstory="You're an expert in research and you love to learn new things.",
-        allow_delegation=False,
-    )
+        math_researcher = Agent(
+            role="Researcher",
+            goal="You research about math.",
+            backstory="You're an expert in research and you love to learn new things.",
+            allow_delegation=False,
+        )

-    task1 = Task(
-        description="Research a topic to teach a kid aged 6 about math.",
-        expected_output="A topic, explanation, angle, and examples.",
-        agent=math_researcher,
-    )
+        task1 = Task(
+            description="Research a topic to teach a kid aged 6 about math.",
+            expected_output="A topic, explanation, angle, and examples.",
+            agent=math_researcher,
+        )

-    crew = Crew(
-        agents=[math_researcher],
-        tasks=[task1],
-        memory=True,
-    )
+        crew = Crew(
+            agents=[math_researcher],
+            tasks=[task1],
+            memory=True,
+        )

-    crew.kickoff()
+        crew.kickoff()

    assert len(events["MemorySaveStartedEvent"]) == 6
    assert len(events["MemorySaveCompletedEvent"]) == 6
--- a/tests/memory/external_memory_test.py
+++ b/tests/memory/external_memory_test.py
@@ -238,20 +238,20 @@ def test_external_memory_search_events(custom_storage, external_memory_with_mock
    events = defaultdict(list)

    external_memory_with_mocked_config.storage = custom_storage
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemoryQueryStartedEvent)
+        def on_search_started(source, event):
+            events["MemoryQueryStartedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryStartedEvent)
-    def on_search_started(source, event):
-        events["MemoryQueryStartedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryCompletedEvent)
+        def on_search_completed(source, event):
+            events["MemoryQueryCompletedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryCompletedEvent)
-    def on_search_completed(source, event):
-        events["MemoryQueryCompletedEvent"].append(event)
-
-    external_memory_with_mocked_config.search(
-        query="test value",
-        limit=3,
-        score_threshold=0.35,
-    )
+        external_memory_with_mocked_config.search(
+            query="test value",
+            limit=3,
+            score_threshold=0.35,
+        )

    assert len(events["MemoryQueryStartedEvent"]) == 1
    assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -281,7 +281,6 @@ def test_external_memory_search_events(custom_storage, external_memory_with_mock
        'query_time_ms': ANY
    }

-    crewai_event_bus._handlers.clear()


 def test_external_memory_save_events(custom_storage, external_memory_with_mocked_config):
@@ -289,19 +288,20 @@ def test_external_memory_save_events(custom_storage, external_memory_with_mocked

    external_memory_with_mocked_config.storage = custom_storage

-    @crewai_event_bus.on(MemorySaveStartedEvent)
-    def on_save_started(source, event):
-        events["MemorySaveStartedEvent"].append(event)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemorySaveStartedEvent)
+        def on_save_started(source, event):
+            events["MemorySaveStartedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveCompletedEvent)
-    def on_save_completed(source, event):
-        events["MemorySaveCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemorySaveCompletedEvent)
+        def on_save_completed(source, event):
+            events["MemorySaveCompletedEvent"].append(event)

-    external_memory_with_mocked_config.save(
-        value="saving value",
-        metadata={"task": "test_task"},
-        agent="test_agent",
-    )
+        external_memory_with_mocked_config.save(
+            value="saving value",
+            metadata={"task": "test_task"},
+            agent="test_agent",
+        )

    assert len(events["MemorySaveStartedEvent"]) == 1
    assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -329,5 +329,3 @@ def test_external_memory_save_events(custom_storage, external_memory_with_mocked
        'agent_role': "test_agent",
        'save_time_ms': ANY
    }
-
-    crewai_event_bus._handlers.clear()
--- a/tests/memory/long_term_memory_test.py
+++ b/tests/memory/long_term_memory_test.py
@@ -20,23 +20,24 @@ def long_term_memory():
 def test_long_term_memory_save_events(long_term_memory):
    events = defaultdict(list)

-    @crewai_event_bus.on(MemorySaveStartedEvent)
-    def on_save_started(source, event):
-        events["MemorySaveStartedEvent"].append(event)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemorySaveStartedEvent)
+        def on_save_started(source, event):
+            events["MemorySaveStartedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveCompletedEvent)
-    def on_save_completed(source, event):
-        events["MemorySaveCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemorySaveCompletedEvent)
+        def on_save_completed(source, event):
+            events["MemorySaveCompletedEvent"].append(event)

-    memory = LongTermMemoryItem(
-        agent="test_agent",
-        task="test_task",
-        expected_output="test_output",
-        datetime="test_datetime",
-        quality=0.5,
-        metadata={"task": "test_task", "quality": 0.5},
-    )
-    long_term_memory.save(memory)
+        memory = LongTermMemoryItem(
+            agent="test_agent",
+            task="test_task",
+            expected_output="test_output",
+            datetime="test_datetime",
+            quality=0.5,
+            metadata={"task": "test_task", "quality": 0.5},
+        )
+        long_term_memory.save(memory)

    assert len(events["MemorySaveStartedEvent"]) == 1
    assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -64,26 +65,25 @@ def test_long_term_memory_save_events(long_term_memory):
        "save_time_ms": ANY,
    }

-    crewai_event_bus._handlers.clear()
-

 def test_long_term_memory_search_events(long_term_memory):
    events = defaultdict(list)

-    @crewai_event_bus.on(MemoryQueryStartedEvent)
-    def on_search_started(source, event):
-        events["MemoryQueryStartedEvent"].append(event)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemoryQueryStartedEvent)
+        def on_search_started(source, event):
+            events["MemoryQueryStartedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryCompletedEvent)
-    def on_search_completed(source, event):
-        events["MemoryQueryCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryCompletedEvent)
+        def on_search_completed(source, event):
+            events["MemoryQueryCompletedEvent"].append(event)

-    test_query = "test query"
+        test_query = "test query"

-    long_term_memory.search(
-        test_query,
-        latest_n=5
-    )
+        long_term_memory.search(
+            test_query,
+            latest_n=5
+        )

    assert len(events["MemoryQueryStartedEvent"]) == 1
    assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -113,7 +113,6 @@ def test_long_term_memory_search_events(long_term_memory):
        'query_time_ms': ANY
    }

-    crewai_event_bus._handlers.clear()

 def test_save_and_search(long_term_memory):
    memory = LongTermMemoryItem(
--- a/tests/memory/short_term_memory_test.py
+++ b/tests/memory/short_term_memory_test.py
@@ -38,20 +38,21 @@ def short_term_memory():
 def test_short_term_memory_search_events(short_term_memory):
    events = defaultdict(list)

-    @crewai_event_bus.on(MemoryQueryStartedEvent)
-    def on_search_started(source, event):
-        events["MemoryQueryStartedEvent"].append(event)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemoryQueryStartedEvent)
+        def on_search_started(source, event):
+            events["MemoryQueryStartedEvent"].append(event)

-    @crewai_event_bus.on(MemoryQueryCompletedEvent)
-    def on_search_completed(source, event):
-        events["MemoryQueryCompletedEvent"].append(event)
+        @crewai_event_bus.on(MemoryQueryCompletedEvent)
+        def on_search_completed(source, event):
+            events["MemoryQueryCompletedEvent"].append(event)

-    # Call the save method
-    short_term_memory.search(
-        query="test value",
-        limit=3,
-        score_threshold=0.35,
-    )
+        # Call the save method
+        short_term_memory.search(
+            query="test value",
+            limit=3,
+            score_threshold=0.35,
+        )

    assert len(events["MemoryQueryStartedEvent"]) == 1
    assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -81,25 +82,23 @@ def test_short_term_memory_search_events(short_term_memory):
        'query_time_ms': ANY
    }

-    crewai_event_bus._handlers.clear()
-

 def test_short_term_memory_save_events(short_term_memory):
    events = defaultdict(list)
+    with crewai_event_bus.scoped_handlers():
+        @crewai_event_bus.on(MemorySaveStartedEvent)
+        def on_save_started(source, event):
+            events["MemorySaveStartedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveStartedEvent)
-    def on_save_started(source, event):
-        events["MemorySaveStartedEvent"].append(event)
+        @crewai_event_bus.on(MemorySaveCompletedEvent)
+        def on_save_completed(source, event):
+            events["MemorySaveCompletedEvent"].append(event)

-    @crewai_event_bus.on(MemorySaveCompletedEvent)
-    def on_save_completed(source, event):
-        events["MemorySaveCompletedEvent"].append(event)
-
-    short_term_memory.save(
-        value="test value",
-        metadata={"task": "test_task"},
-        agent="test_agent",
-    )
+        short_term_memory.save(
+            value="test value",
+            metadata={"task": "test_task"},
+            agent="test_agent",
+        )

    assert len(events["MemorySaveStartedEvent"]) == 1
    assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -128,8 +127,6 @@ def test_short_term_memory_save_events(short_term_memory):
        'save_time_ms': ANY
    }

-    crewai_event_bus._handlers.clear()
-
 def test_save_and_search(short_term_memory):
    memory = ShortTermMemoryItem(
        data="""test value test value test value test value test value test value