From 174215980b2d8eadbf461020779b501d538b5c74 Mon Sep 17 00:00:00 2001 From: Lucas Gomide Date: Tue, 1 Jul 2025 23:07:56 -0300 Subject: [PATCH] test: use scoped_handlers while listener Memory events --- tests/crew_test.py | 83 +++++++++++++------------- tests/memory/external_memory_test.py | 50 ++++++++-------- tests/memory/long_term_memory_test.py | 57 +++++++++--------- tests/memory/short_term_memory_test.py | 53 ++++++++-------- 4 files changed, 119 insertions(+), 124 deletions(-) diff --git a/tests/crew_test.py b/tests/crew_test.py index 31cfbd96d..88e3fa4f6 100644 --- 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 diff --git a/tests/memory/external_memory_test.py b/tests/memory/external_memory_test.py index f072aefaa..c37b5eb85 100644 --- 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() diff --git a/tests/memory/long_term_memory_test.py b/tests/memory/long_term_memory_test.py index ae8c00fd4..0628a3851 100644 --- 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( diff --git a/tests/memory/short_term_memory_test.py b/tests/memory/short_term_memory_test.py index 47c427736..020b3a94c 100644 --- 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