mirror of https://github.com/aliasrobotics/cai.git
433 lines
15 KiB
Python
433 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Backward compatibility tests for TUI streaming.
|
|
|
|
This test suite ensures that the streaming implementation:
|
|
- Doesn't break non-streaming mode
|
|
- Works with legacy display methods
|
|
- Maintains compatibility with existing CLI commands
|
|
- Preserves all existing display features
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from typing import Dict, List, Optional, Any
|
|
from unittest.mock import MagicMock, patch, Mock
|
|
import json
|
|
|
|
# Add src to path
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
|
|
|
|
from cai.tui.display.streaming_display import StreamingDisplay
|
|
from cai.tui.display.tool_display import ToolDisplay
|
|
from cai.tui.display.agent_display import AgentDisplay
|
|
from cai.tui.display.manager import DisplayManager
|
|
from cai.tui.display.base import DisplayContext, OutputType
|
|
from cai.tui.display.panel_formatter import PanelFormatter
|
|
from rich.panel import Panel
|
|
from rich.console import Console
|
|
|
|
|
|
class TestStreamingBackwardCompatibility(unittest.TestCase):
|
|
"""Test backward compatibility of streaming implementation."""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment."""
|
|
# Test both streaming and non-streaming modes
|
|
self.original_stream_env = os.environ.get("CAI_STREAM", "false")
|
|
|
|
# Create display components
|
|
self.streaming_display = StreamingDisplay()
|
|
self.tool_display = ToolDisplay()
|
|
self.agent_display = AgentDisplay()
|
|
self.display_manager = DisplayManager()
|
|
|
|
# Mock terminal output
|
|
self.mock_terminal = MagicMock()
|
|
self.mock_terminal.write = MagicMock()
|
|
self.mock_terminal.print = MagicMock()
|
|
|
|
# Patch get_terminal_output
|
|
self.patcher = patch('cai.tui.core.terminal_console.get_terminal_output')
|
|
self.mock_get_terminal = self.patcher.start()
|
|
self.mock_get_terminal.return_value = self.mock_terminal
|
|
|
|
def tearDown(self):
|
|
"""Clean up."""
|
|
self.patcher.stop()
|
|
os.environ["CAI_STREAM"] = self.original_stream_env
|
|
|
|
def test_non_streaming_mode_compatibility(self):
|
|
"""Test that non-streaming mode still works."""
|
|
os.environ["CAI_STREAM"] = "false"
|
|
|
|
context = DisplayContext(
|
|
terminal_id="test-1",
|
|
terminal_number=1,
|
|
agent_name="Test Agent",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test direct display methods (non-streaming)
|
|
# Agent message
|
|
self.agent_display.display(context, {
|
|
"messages": [{
|
|
"role": "assistant",
|
|
"content": "This is a non-streaming response"
|
|
}],
|
|
"model": "gpt-4"
|
|
})
|
|
|
|
# Tool output
|
|
self.tool_display.display(context, {
|
|
"tool_name": "generic_linux_command",
|
|
"args": {"command": "echo", "args": "test"},
|
|
"output": "test",
|
|
"execution_info": {"status": "completed"}
|
|
})
|
|
|
|
# Verify output was written
|
|
self.assertTrue(self.mock_terminal.write.called)
|
|
self.assertEqual(self.mock_terminal.write.call_count, 2)
|
|
|
|
def test_legacy_display_methods(self):
|
|
"""Test that legacy display methods still work."""
|
|
context = DisplayContext(
|
|
terminal_id="test-2",
|
|
terminal_number=1,
|
|
agent_name="Legacy Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test direct panel creation (legacy method)
|
|
panel = PanelFormatter.create_agent_panel(
|
|
agent_name="Legacy Agent",
|
|
message="Testing legacy display",
|
|
metadata={"model": "gpt-3.5-turbo"},
|
|
streaming=False
|
|
)
|
|
|
|
# Should be able to write panel directly
|
|
self.mock_terminal.write(panel)
|
|
|
|
# Verify panel was created and written
|
|
self.assertIsInstance(panel, Panel)
|
|
self.assertTrue(self.mock_terminal.write.called)
|
|
|
|
def test_display_manager_routing(self):
|
|
"""Test that display manager correctly routes events."""
|
|
context = DisplayContext(
|
|
terminal_id="test-3",
|
|
terminal_number=1,
|
|
agent_name="Router Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test various event types
|
|
events = [
|
|
('agent_message', {
|
|
"messages": [{"role": "assistant", "content": "Test message"}]
|
|
}),
|
|
('tool_output', {
|
|
"tool_name": "test_tool",
|
|
"output": "Tool output"
|
|
}),
|
|
('stream_start', {
|
|
"stream_id": "test-stream",
|
|
"content_type": "text"
|
|
}),
|
|
('stream_update', {
|
|
"stream_id": "test-stream",
|
|
"content": "Streaming content"
|
|
}),
|
|
('stream_finish', {
|
|
"stream_id": "test-stream"
|
|
})
|
|
]
|
|
|
|
for event_type, data in events:
|
|
# Should not raise any exceptions
|
|
self.display_manager.dispatch(event_type, context, data)
|
|
|
|
def test_panel_formatter_compatibility(self):
|
|
"""Test that panel formatter methods remain compatible."""
|
|
# Test all panel creation methods
|
|
|
|
# Agent panel
|
|
agent_panel = PanelFormatter.create_agent_panel(
|
|
agent_name="Test Agent",
|
|
message="Test message",
|
|
metadata={"timestamp": "10:00:00"},
|
|
streaming=True
|
|
)
|
|
self.assertIsInstance(agent_panel, Panel)
|
|
|
|
# Tool panel
|
|
tool_panel = PanelFormatter.create_tool_panel(
|
|
tool_name="generic_linux_command",
|
|
args={"command": "ls"},
|
|
output="file1.txt\nfile2.txt",
|
|
execution_info={"tool_time": 0.5},
|
|
token_info={},
|
|
streaming=False
|
|
)
|
|
self.assertIsInstance(tool_panel, Panel)
|
|
|
|
# Thinking panel
|
|
thinking_panel = PanelFormatter.create_thinking_panel(
|
|
agent_name="Thinking Agent",
|
|
thinking_content="Processing request...",
|
|
model_name="o1-preview",
|
|
finished=True
|
|
)
|
|
self.assertIsInstance(thinking_panel, Panel)
|
|
|
|
# Error panel
|
|
error_panel = PanelFormatter.create_error_panel(
|
|
error_type="ValueError",
|
|
error_message="Invalid input",
|
|
details="Expected string, got int"
|
|
)
|
|
self.assertIsInstance(error_panel, Panel)
|
|
|
|
def test_mixed_streaming_and_non_streaming(self):
|
|
"""Test mixing streaming and non-streaming display."""
|
|
context = DisplayContext(
|
|
terminal_id="test-4",
|
|
terminal_number=1,
|
|
agent_name="Mixed Mode",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Non-streaming tool output
|
|
self.tool_display.display(context, {
|
|
"tool_name": "read_file",
|
|
"args": {"path": "test.txt"},
|
|
"output": "File contents",
|
|
"execution_info": {"status": "completed"}
|
|
})
|
|
|
|
# Streaming agent response
|
|
stream_id = "mixed-stream"
|
|
self.streaming_display.start_streaming(context, stream_id, {
|
|
"content_type": "text",
|
|
"model": "gpt-4"
|
|
})
|
|
|
|
self.streaming_display.update_streaming(stream_id, {
|
|
"content": "This is a streaming response"
|
|
})
|
|
|
|
self.streaming_display.finish_streaming(stream_id, {
|
|
"final_stats": {"output_tokens": 5}
|
|
})
|
|
|
|
# Non-streaming agent message
|
|
self.agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Final message"}]
|
|
})
|
|
|
|
# All displays should work without conflicts
|
|
self.assertGreaterEqual(self.mock_terminal.write.call_count, 2)
|
|
|
|
def test_context_history_preservation(self):
|
|
"""Test that context history is preserved with streaming."""
|
|
context = DisplayContext(
|
|
terminal_id="test-5",
|
|
terminal_number=1,
|
|
agent_name="History Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Add outputs of different types
|
|
context.add_output({
|
|
"type": OutputType.AGENT_MESSAGE,
|
|
"content": "First message"
|
|
})
|
|
|
|
context.add_output({
|
|
"type": OutputType.TOOL_OUTPUT,
|
|
"tool_name": "test_tool",
|
|
"output": "Tool result"
|
|
})
|
|
|
|
context.add_output({
|
|
"type": OutputType.STREAMING,
|
|
"content": "Streamed content"
|
|
})
|
|
|
|
# Verify history is maintained
|
|
self.assertEqual(len(context.outputs), 3)
|
|
self.assertEqual(context.outputs[0]["type"], OutputType.AGENT_MESSAGE)
|
|
self.assertEqual(context.outputs[1]["type"], OutputType.TOOL_OUTPUT)
|
|
self.assertEqual(context.outputs[2]["type"], OutputType.STREAMING)
|
|
|
|
def test_terminal_output_methods(self):
|
|
"""Test both write and print methods on terminal output."""
|
|
context = DisplayContext(
|
|
terminal_id="test-6",
|
|
terminal_number=1,
|
|
agent_name="Terminal Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test with terminal that has write method
|
|
terminal_with_write = MagicMock()
|
|
terminal_with_write.write = MagicMock()
|
|
|
|
with patch('cai.tui.core.terminal_console.get_terminal_output', return_value=terminal_with_write):
|
|
self.agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Test write method"}]
|
|
})
|
|
|
|
self.assertTrue(terminal_with_write.write.called)
|
|
|
|
# Test with terminal that has print method
|
|
terminal_with_print = MagicMock()
|
|
terminal_with_print.print = MagicMock()
|
|
del terminal_with_print.write # Ensure no write method
|
|
|
|
with patch('cai.tui.core.terminal_console.get_terminal_output', return_value=terminal_with_print):
|
|
# Create new display instance to test print path
|
|
agent_display = AgentDisplay()
|
|
agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Test print method"}]
|
|
})
|
|
|
|
# Should fall back to console methods
|
|
|
|
def test_streaming_api_presence(self):
|
|
"""Test that streaming API doesn't break terminals without it."""
|
|
context = DisplayContext(
|
|
terminal_id="test-7",
|
|
terminal_number=1,
|
|
agent_name="API Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Terminal without streaming API
|
|
basic_terminal = MagicMock()
|
|
basic_terminal.write = MagicMock()
|
|
# No streaming methods
|
|
|
|
with patch('cai.tui.core.terminal_console.get_terminal_output', return_value=basic_terminal):
|
|
# Should gracefully handle missing streaming API
|
|
stream_id = "no-api-stream"
|
|
self.streaming_display.start_streaming(context, stream_id, {
|
|
"content_type": "text"
|
|
})
|
|
|
|
self.streaming_display.update_streaming(stream_id, {
|
|
"content": "Content without streaming API"
|
|
})
|
|
|
|
self.streaming_display.finish_streaming(stream_id, {})
|
|
|
|
# Should fall back to regular write
|
|
self.assertTrue(basic_terminal.write.called)
|
|
|
|
def test_deduplication_still_works(self):
|
|
"""Test that deduplication logic is preserved."""
|
|
context = DisplayContext(
|
|
terminal_id="test-8",
|
|
terminal_number=1,
|
|
agent_name="Dedup Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Display same tool output multiple times
|
|
tool_data = {
|
|
"tool_name": "echo",
|
|
"args": {"text": "duplicate"},
|
|
"output": "duplicate",
|
|
"call_id": "same_call_id"
|
|
}
|
|
|
|
# First display
|
|
self.tool_display.display(context, tool_data)
|
|
first_call_count = self.mock_terminal.write.call_count
|
|
|
|
# Immediate duplicate should be skipped
|
|
self.tool_display.display(context, tool_data)
|
|
self.assertEqual(self.mock_terminal.write.call_count, first_call_count)
|
|
|
|
# After time passes, should display again
|
|
import time
|
|
time.sleep(0.6) # Wait more than dedup window
|
|
self.tool_display.display(context, tool_data)
|
|
self.assertGreater(self.mock_terminal.write.call_count, first_call_count)
|
|
|
|
def test_error_handling_compatibility(self):
|
|
"""Test error handling remains compatible."""
|
|
context = DisplayContext(
|
|
terminal_id="test-9",
|
|
terminal_number=1,
|
|
agent_name="Error Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test with various error conditions
|
|
# Missing required data
|
|
self.agent_display.display(context, {}) # No messages
|
|
|
|
# Invalid data types
|
|
self.tool_display.display(context, {
|
|
"tool_name": None,
|
|
"output": 123 # Should be string
|
|
})
|
|
|
|
# Streaming with missing stream_id
|
|
self.streaming_display.update_streaming("non_existent_stream", {
|
|
"content": "Update to non-existent stream"
|
|
})
|
|
|
|
# None of these should raise exceptions
|
|
# System should handle gracefully
|
|
|
|
def test_token_stats_compatibility(self):
|
|
"""Test token statistics display compatibility."""
|
|
context = DisplayContext(
|
|
terminal_id="test-10",
|
|
terminal_number=1,
|
|
agent_name="Stats Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Test with full token info
|
|
full_token_info = {
|
|
"interaction_input_tokens": 100,
|
|
"interaction_output_tokens": 50,
|
|
"interaction_cost": 0.0025,
|
|
"session_total_cost": 0.0175,
|
|
"context_usage_pct": 15.5
|
|
}
|
|
|
|
self.agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Response with full stats"}],
|
|
"token_info": full_token_info
|
|
})
|
|
|
|
# Test with partial token info
|
|
partial_token_info = {
|
|
"interaction_input_tokens": 75,
|
|
"interaction_output_tokens": 25
|
|
# Missing costs and context usage
|
|
}
|
|
|
|
self.agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Response with partial stats"}],
|
|
"token_info": partial_token_info
|
|
})
|
|
|
|
# Test with no token info
|
|
self.agent_display.display(context, {
|
|
"messages": [{"role": "assistant", "content": "Response without stats"}]
|
|
})
|
|
|
|
# All should display without errors
|
|
self.assertGreaterEqual(self.mock_terminal.write.call_count, 3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2) |