import os import json import pytest from unittest.mock import patch, Mock, MagicMock from typing import Any, cast, Dict, List, Union from pathlib import Path from autogenstudio.lite import LiteStudio @pytest.fixture def sample_team_file(tmp_path: Path) -> str: """Fixture for creating a sample team JSON file""" team_data: Dict[str, Union[str, List[Any]]] = { "type": "SampleTeam", "name": "Test Team", "description": "A test team for lite mode", "agents": [] } team_file = tmp_path / "test_team.json" with open(team_file, 'w') as f: json.dump(team_data, f) return str(team_file) @pytest.fixture def sample_team_object() -> Dict[str, Union[str, List[Any]]]: """Fixture for sample team object""" return { "type": "TestTeam", "name": "Object Team", "agents": [] } def load_team_from_file(file_path: str) -> Dict[str, Any]: """Helper function to load team data from file""" with open(file_path, 'r') as f: return json.load(f) def test_init_with_file_team(sample_team_file: str) -> None: """Test LiteStudio initialization with a team file""" studio = LiteStudio( team=sample_team_file, host="localhost", port=9090, session_name="Test Session" ) assert studio.host == "localhost" assert studio.port == 9090 assert studio.session_name == "Test Session" assert studio.auto_open is True assert studio.team_file_path == os.path.abspath(sample_team_file) assert studio.server_process is None def test_init_with_team_object(sample_team_object: Dict[str, Union[str, List[Any]]]) -> None: """Test LiteStudio initialization with a team object""" studio = LiteStudio(team=sample_team_object, port=9091) # Should create a temporary file assert studio.team_file_path is not None assert os.path.exists(studio.team_file_path) assert studio.port == 9091 # Verify content matches team_data = load_team_from_file(studio.team_file_path) assert team_data == sample_team_object @patch('autogenstudio.gallery.builder.create_default_lite_team') def test_init_with_no_team(mock_create_default: MagicMock) -> None: """Test LiteStudio initialization with no team (should create default)""" mock_create_default.return_value = "/tmp/default_team.json" with patch('os.path.exists', return_value=True): studio = LiteStudio() mock_create_default.assert_called_once() assert studio.team_file_path is not None # Verify default parameters assert studio.host == "127.0.0.1" assert studio.port == 8080 assert studio.auto_open is True assert studio.session_name == "Lite Session" def test_init_with_invalid_file() -> None: """Test LiteStudio initialization with invalid team file""" with pytest.raises(FileNotFoundError, match="Team file not found"): LiteStudio(team="/nonexistent/file.json") def test_setup_environment(sample_team_file: str) -> None: """Test environment variable setup""" studio = LiteStudio(team=sample_team_file, host="127.0.0.1", port=8080) env_file_path = studio._setup_environment() # type: ignore # Read the environment file to verify contents env_vars = {} with open(env_file_path, 'r') as f: for line in f: if '=' in line: key, value = line.strip().split('=', 1) env_vars[key] = value expected_vars = { "AUTOGENSTUDIO_HOST": "127.0.0.1", "AUTOGENSTUDIO_PORT": "8080", "AUTOGENSTUDIO_LITE_MODE": "true", "AUTOGENSTUDIO_API_DOCS": "false", "AUTOGENSTUDIO_AUTH_DISABLED": "true", "AUTOGENSTUDIO_DATABASE_URI": "sqlite:///:memory:", } for key, value in expected_vars.items(): assert key in env_vars assert env_vars[key] == value @patch('uvicorn.run') @patch('threading.Thread') @patch('webbrowser.open') def test_start_foreground(mock_browser: MagicMock, mock_thread: MagicMock, mock_uvicorn: MagicMock, sample_team_file: str) -> None: """Test starting studio in foreground mode""" studio = LiteStudio(team=sample_team_file, auto_open=True) studio.start(background=False) # Should call uvicorn.run mock_uvicorn.assert_called_once() # Should setup browser opening thread mock_thread.assert_called_once() @patch('uvicorn.run') @patch('threading.Thread') def test_start_background(mock_thread_class: MagicMock, mock_uvicorn: MagicMock, sample_team_file: str) -> None: """Test starting studio in background mode""" studio = LiteStudio(team=sample_team_file) # Mock the server thread mock_server_thread = Mock() mock_thread_class.return_value = mock_server_thread studio.start(background=True) # Should create and start server thread mock_thread_class.assert_called() # Note: start() is called twice - once for browser opening, once for server assert mock_server_thread.start.call_count >= 1 assert studio.server_thread == mock_server_thread def test_context_manager(sample_team_file: str) -> None: """Test LiteStudio as context manager""" with patch.object(LiteStudio, 'start') as mock_start: with patch.object(LiteStudio, 'stop') as mock_stop: with LiteStudio(team=sample_team_file) as studio: assert studio is not None mock_start.assert_called_once_with(background=True) mock_stop.assert_called_once() def test_stop_with_server_thread(sample_team_file: str) -> None: """Test stopping studio with active server thread""" studio = LiteStudio(team=sample_team_file) # Mock server thread mock_thread = Mock() mock_thread.is_alive.return_value = True studio.server_thread = mock_thread studio.stop() mock_thread.join.assert_called_once_with(timeout=5) @patch('subprocess.run') def test_shutdown_port(mock_subprocess: MagicMock) -> None: """Test shutting down a specific port""" # Mock subprocess return with stdout containing PIDs mock_result = Mock() mock_result.returncode = 0 mock_result.stdout = "1234\n5678" mock_subprocess.return_value = mock_result LiteStudio.shutdown_port(8080) # Should attempt to find and kill process on port mock_subprocess.assert_called() @patch('uvicorn.run') def test_start_twice_raises_error(mock_uvicorn: MagicMock, sample_team_file: str) -> None: """Test that starting an already running studio raises an error""" studio = LiteStudio(team=sample_team_file) # Mock that server is already running studio.server_thread = Mock() studio.server_thread.is_alive.return_value = True with pytest.raises(RuntimeError, match="already running"): studio.start() def test_init_with_team_object_with_serialization_methods() -> None: """Test LiteStudio initialization with objects that have serialization methods""" # Mock object with dump_component method (like AutoGen teams) class MockTeamWithDumpComponent: def dump_component(self) -> Any: class MockComponent: def model_dump(self) -> Dict[str, Union[str, List[Any]]]: return {"type": "MockTeam", "name": "Serialized Team", "participants": []} return MockComponent() mock_team = MockTeamWithDumpComponent() # Cast to Any since the runtime handles this properly studio = LiteStudio(team=cast(Any, mock_team), port=9092) # Should create a temporary file with serialized content assert studio.team_file_path is not None assert os.path.exists(studio.team_file_path) # Verify content matches serialization team_data = load_team_from_file(studio.team_file_path) assert team_data["type"] == "MockTeam" assert team_data["name"] == "Serialized Team" def test_init_with_team_object_model_dump() -> None: """Test LiteStudio initialization with Pydantic v2 style objects""" class MockPydanticTeam: def model_dump(self): return {"type": "PydanticTeam", "name": "Model Dump Team"} mock_team = MockPydanticTeam() # Cast to Any since the runtime handles this properly studio = LiteStudio(team=cast(Any, mock_team), port=9093) # Verify content team_data = load_team_from_file(studio.team_file_path) assert team_data["type"] == "PydanticTeam" assert team_data["name"] == "Model Dump Team" def test_init_with_unsupported_team_object() -> None: """Test LiteStudio initialization with unsupported team object""" class UnsupportedTeam: def some_other_method(self): return "not serializable" with pytest.raises(ValueError, match="Cannot serialize team object"): # Cast to Any since we're intentionally testing unsupported type LiteStudio(team=cast(Any, UnsupportedTeam())) def test_init_with_path_object(tmp_path: Path) -> None: """Test LiteStudio initialization with Path object""" team_data: Dict[str, Union[str, List[Any]]] = { "type": "PathTeam", "name": "Path Test Team", "agents": [] } team_file = tmp_path / "path_team.json" with open(team_file, 'w') as f: json.dump(team_data, f) # Use Path object instead of string studio = LiteStudio(team=team_file, port=9094) assert studio.team_file_path == str(team_file.absolute()) assert os.path.exists(studio.team_file_path) # Verify content team_data_loaded = load_team_from_file(studio.team_file_path) assert team_data_loaded == team_data def test_init_with_component_model() -> None: """Test LiteStudio initialization with ComponentModel""" # Since ComponentModel is more complex to create directly, # we'll test this by mocking it in the _load_team method team_dict: Dict[str, Union[str, List[Any]]] = { "type": "ComponentTeam", "name": "Component Model Team", "participants": [] } # Create a mock that behaves like ComponentModel from unittest.mock import Mock mock_component = Mock() mock_component.model_dump.return_value = team_dict # Test that it gets handled in the ComponentModel branch studio = LiteStudio(team=team_dict, port=9095) # Use dict instead for now # Verify the dict path works (ComponentModel would use same serialization) team_data = load_team_from_file(studio.team_file_path) assert team_data["type"] == "ComponentTeam" assert team_data["name"] == "Component Model Team"