mirror of
https://github.com/microsoft/autogen.git
synced 2025-08-04 14:52:10 +00:00
b2cef7f47c
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
312 lines
10 KiB
Python
312 lines
10 KiB
Python
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"
|