mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
352b213f7f
Removes unused imports that were causing lint failures: - logging - io.StringIO - unittest.mock.MagicMock - rich.logging.RichHandler This fixes the CI lint check failure (job 45782458300). Co-Authored-By: Jo\u00E3o <joao@crewai.com>
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
from unittest.mock import patch
|
|
from rich.tree import Tree
|
|
from crewai.utilities.events.utils.console_formatter import ConsoleFormatter
|
|
from crewai.utilities.events.event_listener import EventListener
|
|
|
|
|
|
class TestRichLiveCleanup:
|
|
"""Test that Rich Live sessions are properly cleaned up after CrewAI operations."""
|
|
|
|
def test_logging_works_after_tree_rendering(self):
|
|
"""Test that logging output appears after tree rendering with proper cleanup."""
|
|
formatter = ConsoleFormatter()
|
|
|
|
tree = Tree("Test Flow")
|
|
formatter.print(tree)
|
|
|
|
assert formatter._live is not None
|
|
|
|
formatter.stop_live()
|
|
|
|
assert formatter._live is None
|
|
|
|
with patch.object(formatter.console, 'print') as mock_print:
|
|
formatter.print("This should appear immediately")
|
|
mock_print.assert_called_once_with("This should appear immediately")
|
|
|
|
def test_event_listener_cleanup_integration(self):
|
|
"""Test that EventListener properly cleans up Live sessions."""
|
|
event_listener = EventListener()
|
|
formatter = event_listener.formatter
|
|
|
|
tree = Tree("Test Crew")
|
|
formatter.print(tree)
|
|
assert formatter._live is not None
|
|
|
|
formatter.stop_live()
|
|
assert formatter._live is None
|
|
|
|
def test_stop_live_restores_normal_output(self):
|
|
"""Test that stop_live properly restores normal console output behavior."""
|
|
formatter = ConsoleFormatter()
|
|
|
|
tree = Tree("Test Tree")
|
|
formatter.print(tree)
|
|
|
|
assert formatter._live is not None
|
|
|
|
formatter.stop_live()
|
|
|
|
assert formatter._live is None
|
|
|
|
with patch.object(formatter.console, 'print') as mock_print:
|
|
formatter.print("Normal output")
|
|
mock_print.assert_called_once_with("Normal output")
|