""" Tests for agency_swarm.agent.attachment_manager module. Tests the AttachmentManager class functionality including vector store management, file attachment handling, and cleanup operations. """ from unittest.mock import Mock import pytest from agents import CodeInterpreterTool from agents.exceptions import AgentsException from agency_swarm.agent.attachment_manager import AttachmentManager class TestAttachmentManager: """Test AttachmentManager class functionality.""" def test_init_without_file_manager(self): """Test AttachmentManager initialization when agent has no file_manager.""" # Create a mock agent without file_manager mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = None # Should raise AgentsException on line 48 with pytest.raises( AgentsException, match="Cannot use AttachmentManager for agent TestAgent without file manager" ): AttachmentManager(mock_agent) def test_init_attachments_vs_existing_vector_store(self): """Test init_attachments_vs when vector store already exists.""" # Setup mock agent with file_manager mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Mock vector store list response with existing VS mock_vs_data = Mock() mock_vs_data.name = "attachments_vs" mock_vs_data.id = "vs_test123" mock_vs_list = Mock() mock_vs_list.data = [mock_vs_data] mock_agent.client_sync.vector_stores.list.return_value = mock_vs_list attachment_manager = AttachmentManager(mock_agent) # Call init_attachments_vs with existing VS name - should return existing ID result = attachment_manager.init_attachments_vs("attachments_vs") assert result == "vs_test123" mock_agent.client_sync.vector_stores.list.assert_called_once() mock_agent.client_sync.vector_stores.create.assert_not_called() def test_init_attachments_vs_create_new(self): """Test init_attachments_vs when creating new vector store.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Mock vector store list response with no existing VS mock_vs_list = Mock() mock_vs_list.data = [] # Mock create response mock_created_vs = Mock() mock_created_vs.id = "vs_new456" mock_agent.client_sync.vector_stores.list.return_value = mock_vs_list mock_agent.client_sync.vector_stores.create.return_value = mock_created_vs attachment_manager = AttachmentManager(mock_agent) # Call init_attachments_vs with new VS name - should create new VS result = attachment_manager.init_attachments_vs("new_vs") assert result == "vs_new456" mock_agent.client_sync.vector_stores.list.assert_called_once() mock_agent.client_sync.vector_stores.create.assert_called_once_with(name="new_vs") @pytest.mark.asyncio async def test_sort_file_attachments_unsupported_extension(self): """Test sort_file_attachments with unsupported file extension.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Mock _get_filename_by_id to return file with unsupported extension attachment_manager = AttachmentManager(mock_agent) attachment_manager._get_filename_by_id = Mock(return_value="test.xyz") # Should raise AgentsException with pytest.raises(AgentsException, match="Unsupported file extension: .xyz for file test.xyz"): await attachment_manager.sort_file_attachments(["file-123"]) @pytest.mark.asyncio async def test_prepare_and_attach_files_rejects_message_files(self): """Test prepare_and_attach_files rejects deprecated message_files usage.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() attachment_manager = AttachmentManager(mock_agent) with pytest.raises(TypeError, match="message_files"): await attachment_manager.prepare_and_attach_files( [{"role": "user", "content": "hello"}], None, {"message_files": ["file-123"]}, ) @pytest.mark.asyncio async def test_prepare_and_attach_files_adds_code_interpreter_instruction(self): """Code-interpreter attachments should tell the model how to inspect uploaded files.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() attachment_manager = AttachmentManager(mock_agent) attachment_manager._get_filename_by_id = Mock(return_value="report.txt") message_items = [{"role": "user", "content": "read the file"}] await attachment_manager.prepare_and_attach_files(message_items, ["file-123"], {}) assert message_items == [ { "role": "user", "content": [ {"type": "input_text", "text": "read the file"}, { "type": "input_text", "_agency_swarm_ephemeral": True, "text": ( "The uploaded file(s) report.txt are available through the code_interpreter tool. " "Use that tool to inspect their contents before answering." ), }, ], } ] mock_agent.file_manager.add_code_interpreter_tool.assert_called_once_with(["file-123"]) def test_attachments_cleanup_code_interpreter_files(self): """Test attachments_cleanup with temporary code interpreter files.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Add a CodeInterpreterTool with temp files mock_code_tool = Mock(spec=CodeInterpreterTool) mock_code_tool.tool_config = {"container": {"file_ids": ["file-123", "file-456", "file-789"]}} mock_agent.tools = [mock_code_tool] attachment_manager = AttachmentManager(mock_agent) attachment_manager._temp_code_interpreter_file_ids = ["file-123", "file-789"] # Call cleanup - should remove temp files from tool attachment_manager.attachments_cleanup() # Verify temp files were removed but tool kept (still has file-456) expected_file_ids = ["file-456"] assert mock_code_tool.tool_config["container"]["file_ids"] == expected_file_ids assert mock_code_tool in mock_agent.tools def test_attachments_cleanup_remove_code_interpreter_tool(self): """Test attachments_cleanup removing CodeInterpreterTool when no files left.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Add a CodeInterpreterTool with only temp files mock_code_tool = Mock(spec=CodeInterpreterTool) mock_code_tool.tool_config = {"container": {"file_ids": ["file-123", "file-456"]}} mock_agent.tools = [mock_code_tool] attachment_manager = AttachmentManager(mock_agent) attachment_manager._temp_code_interpreter_file_ids = ["file-123", "file-456"] # Call cleanup - should remove entire tool attachment_manager.attachments_cleanup() # Verify tool was completely removed assert mock_code_tool not in mock_agent.tools def test_attachments_cleanup_code_interpreter_string_container(self): """Test attachments_cleanup with CodeInterpreterTool using string container.""" mock_agent = Mock() mock_agent.name = "TestAgent" mock_agent.file_manager = Mock() # Add a CodeInterpreterTool with string container (can't modify) mock_code_tool = Mock(spec=CodeInterpreterTool) mock_code_tool.tool_config = {"container": "some_container_id"} mock_agent.tools = [mock_code_tool] attachment_manager = AttachmentManager(mock_agent) attachment_manager._temp_code_interpreter_file_ids = ["file-123"] # Call cleanup - should handle string container gracefully attachment_manager.attachments_cleanup() # Should not raise exception # Verify tool configuration wasn't modified assert mock_code_tool.tool_config["container"] == "some_container_id"