aixtools 0.1.10__py3-none-any.whl → 0.2.0__py3-none-any.whl

tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py

@@ -1,156 +0,0 @@
-"""Tests for the Pydantic AI adapter storage module."""
-
-import unittest
-from unittest.mock import MagicMock
-
-from aixtools.a2a.google_sdk.pydantic_ai_adapter.storage import (
-    PydanticAiAgentHistoryStorage,
-    InMemoryHistoryStorage,
-)
-
-
-class TestInMemoryHistoryStorage(unittest.TestCase):
-    """Tests for the InMemoryHistoryStorage class."""
-
-    def setUp(self):
-        self.storage = InMemoryHistoryStorage()
-
-    def test_init(self):
-        """Test InMemoryHistoryStorage initialization."""
-        self.assertEqual(self.storage.storage, {})
-
-    def test_get_nonexistent_task(self):
-        """Test getting history for a task that doesn't exist."""
-        result = self.storage.get("nonexistent_task")
-        self.assertIsNone(result)
-
-    def test_store_and_get_messages(self):
-        """Test storing and retrieving messages."""
-        task_id = "test_task_1"
-        # Use simple mock objects that can be stored
-        messages = [MagicMock(), MagicMock()]
-        
-        # Store the messages
-        self.storage.store(task_id, messages)
-        
-        # Retrieve the messages
-        result = self.storage.get(task_id)
-        
-        self.assertIsNotNone(result)
-        self.assertEqual(result, messages)
-        self.assertEqual(len(result), 2)
-
-    def test_store_overwrites_existing(self):
-        """Test that storing overwrites existing messages for the same task."""
-        task_id = "test_task_3"
-        
-        # Store initial messages
-        initial_messages = [MagicMock()]
-        self.storage.store(task_id, initial_messages)
-        
-        # Store new messages (should overwrite)
-        new_messages = [MagicMock(), MagicMock()]
-        self.storage.store(task_id, new_messages)
-        
-        # Retrieve and verify new messages are stored
-        result = self.storage.get(task_id)
-        
-        self.assertIsNotNone(result)
-        self.assertEqual(result, new_messages)
-        self.assertEqual(len(result), 2)
-        self.assertNotEqual(result, initial_messages)
-
-    def test_multiple_tasks(self):
-        """Test storing and retrieving messages for multiple tasks."""
-        task1_id = "task_1"
-        task2_id = "task_2"
-        
-        task1_messages = [MagicMock()]
-        task2_messages = [MagicMock(), MagicMock()]
-        
-        # Store messages for both tasks
-        self.storage.store(task1_id, task1_messages)
-        self.storage.store(task2_id, task2_messages)
-        
-        # Retrieve and verify both tasks' messages
-        result1 = self.storage.get(task1_id)
-        result2 = self.storage.get(task2_id)
-        
-        self.assertIsNotNone(result1)
-        self.assertIsNotNone(result2)
-        self.assertEqual(result1, task1_messages)
-        self.assertEqual(result2, task2_messages)
-        self.assertNotEqual(result1, result2)
-
-    def test_store_empty_list(self):
-        """Test storing an empty list of messages."""
-        task_id = "empty_task"
-        empty_messages = []
-        
-        self.storage.store(task_id, empty_messages)
-        result = self.storage.get(task_id)
-        
-        self.assertIsNotNone(result)
-        self.assertEqual(result, empty_messages)
-        self.assertEqual(len(result), 0)
-
-    def test_get_after_multiple_stores(self):
-        """Test that get returns the most recent store for a task."""
-        task_id = "update_task"
-        
-        # Store multiple times
-        messages1 = [MagicMock()]
-        messages2 = [MagicMock()]
-        messages3 = [MagicMock(), MagicMock()]
-        
-        self.storage.store(task_id, messages1)
-        self.storage.store(task_id, messages2)
-        self.storage.store(task_id, messages3)
-        
-        result = self.storage.get(task_id)
-        
-        self.assertIsNotNone(result)
-        self.assertEqual(result, messages3)
-
-    def test_storage_isolation(self):
-        """Test that different storage instances are isolated."""
-        storage1 = InMemoryHistoryStorage()
-        storage2 = InMemoryHistoryStorage()
-        
-        task_id = "isolation_test"
-        messages1 = [MagicMock()]
-        messages2 = [MagicMock()]
-        
-        storage1.store(task_id, messages1)
-        storage2.store(task_id, messages2)
-        
-        result1 = storage1.get(task_id)
-        result2 = storage2.get(task_id)
-        
-        self.assertIsNotNone(result1)
-        self.assertIsNotNone(result2)
-        self.assertEqual(result1, messages1)
-        self.assertEqual(result2, messages2)
-        self.assertNotEqual(result1, result2)
-
-
-class TestPydanticAiAgentHistoryStorageInterface(unittest.TestCase):
-    """Tests for the PydanticAiAgentHistoryStorage abstract interface."""
-
-    def test_cannot_instantiate_abstract_class(self):
-        """Test that the abstract base class cannot be instantiated."""
-        with self.assertRaises(TypeError):
-            PydanticAiAgentHistoryStorage()
-
-    def test_inmemory_implements_interface(self):
-        """Test that InMemoryHistoryStorage properly implements the interface."""
-        storage = InMemoryHistoryStorage()
-        
-        # Verify it's an instance of the abstract base class
-        self.assertIsInstance(storage, PydanticAiAgentHistoryStorage)
-        
-        # Verify it has the required methods
-        self.assertTrue(hasattr(storage, 'get'))
-        self.assertTrue(hasattr(storage, 'store'))
-        self.assertTrue(callable(getattr(storage, 'get')))
-        self.assertTrue(callable(getattr(storage, 'store')))

tests/unit/a2a/google_sdk/test_card.py

@@ -1,114 +0,0 @@
-"""Tests for the A2A card module."""
-
-import unittest
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import httpx
-from a2a.client import A2ACardResolver
-from a2a.types import AgentCard
-
-from aixtools.a2a.google_sdk.card import get_agent_card
-
-
-class TestCard(unittest.IsolatedAsyncioTestCase):
-    """Tests for the A2A card module."""
-
-    def setUp(self):
-        self.test_agent_host = "http://localhost:9999"
-
-    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
-    async def test_get_agent_card_success(self, mock_resolver_class):
-        """Test successful retrieval of agent card."""
-        # Setup
-        mock_client = AsyncMock(spec=httpx.AsyncClient)
-        
-        mock_resolver = AsyncMock(spec=A2ACardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        mock_card = MagicMock(spec=AgentCard)
-        mock_card.model_dump_json.return_value = '{"test": "data"}'
-        mock_resolver.get_agent_card.return_value = mock_card
-        
-        # Call the function
-        result = await get_agent_card(mock_client, self.test_agent_host)
-        
-        # Verify the result
-        self.assertEqual(result, mock_card)
-        
-        # Verify the resolver was created correctly
-        mock_resolver_class.assert_called_once_with(
-            httpx_client=mock_client,
-            base_url=self.test_agent_host
-        )
-        
-        # Verify the card was fetched
-        mock_resolver.get_agent_card.assert_called_once()
-        
-        # Verify the URL was set
-        self.assertEqual(result.url, self.test_agent_host)
-
-    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
-    async def test_get_agent_card_failure(self, mock_resolver_class):
-        """Test handling of errors when retrieving agent card."""
-        # Setup
-        mock_client = AsyncMock(spec=httpx.AsyncClient)
-        
-        mock_resolver = AsyncMock(spec=A2ACardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        # Make the resolver raise an exception
-        mock_resolver.get_agent_card.side_effect = Exception("Failed to fetch card")
-        
-        # Call the function and expect an exception
-        with self.assertRaises(RuntimeError) as context:
-            await get_agent_card(mock_client, self.test_agent_host)
-        
-        self.assertIn("Failed to fetch the public agent card", str(context.exception))
-
-    @patch("aixtools.a2a.google_sdk.card.logger")
-    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
-    async def test_get_agent_card_logging(self, mock_resolver_class, mock_logger):
-        """Test that proper logging occurs during card retrieval."""
-        # Setup
-        mock_client = AsyncMock(spec=httpx.AsyncClient)
-        
-        mock_resolver = AsyncMock(spec=A2ACardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        mock_card = MagicMock(spec=AgentCard)
-        mock_card.model_dump_json.return_value = '{"test": "data"}'
-        mock_resolver.get_agent_card.return_value = mock_card
-        
-        # Call the function
-        await get_agent_card(mock_client, self.test_agent_host)
-        
-        # Verify logging calls
-        mock_logger.info.assert_called()
-        self.assertEqual(mock_logger.info.call_count, 2)  # Two info calls in the function
-
-    @patch("aixtools.a2a.google_sdk.card.logger")
-    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
-    async def test_get_agent_card_error_logging(self, mock_resolver_class, mock_logger):
-        """Test that errors are properly logged."""
-        # Setup
-        mock_client = AsyncMock(spec=httpx.AsyncClient)
-        
-        mock_resolver = AsyncMock(spec=A2ACardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        test_error = Exception("Test error")
-        mock_resolver.get_agent_card.side_effect = test_error
-        
-        # Call the function and expect an exception
-        with self.assertRaises(RuntimeError):
-            await get_agent_card(mock_client, self.test_agent_host)
-        
-        # Verify error logging
-        mock_logger.error.assert_called_once()
-        args, kwargs = mock_logger.error.call_args
-        self.assertIn("Critical error fetching public agent card", args[0])
-        self.assertTrue(kwargs.get('exc_info'))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/unit/a2a/google_sdk/test_remote_agent_connection.py

@@ -1,413 +0,0 @@
-"""Tests for the remote agent connection module."""
-
-import unittest
-from unittest.mock import AsyncMock, MagicMock, patch, call
-
-from a2a.client import Client
-from a2a.types import AgentCard, Message, Task, TaskState, TaskStatus, TaskQueryParams
-
-from aixtools.a2a.google_sdk.remote_agent_connection import (
-    RemoteAgentConnection,
-    is_in_terminal_state,
-    is_in_terminal_or_interrupted_state,
-)
-
-
-class TestRemoteAgentConnection(unittest.IsolatedAsyncioTestCase):
-    """Tests for the RemoteAgentConnection class."""
-
-    def setUp(self):
-        self.mock_card = MagicMock(spec=AgentCard)
-        self.mock_client = AsyncMock(spec=Client)
-        self.connection = RemoteAgentConnection(self.mock_card, self.mock_client)
-
-    def test_get_agent_card(self):
-        """Test that get_agent_card returns the stored card."""
-        result = self.connection.get_agent_card()
-        self.assertEqual(result, self.mock_card)
-
-    async def test_send_message_returns_message(self):
-        """Test send_message when it receives a Message response."""
-        mock_message = MagicMock(spec=Message)
-        mock_task = MagicMock(spec=Task)
-        
-        # Mock the async generator to yield a message
-        async def mock_generator():
-            yield mock_message
-            yield (mock_task,)
-        
-        self.mock_client.send_message.return_value = mock_generator()
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message(test_message)
-        
-        self.assertEqual(result, mock_message)
-        self.mock_client.send_message.assert_called_once_with(test_message)
-
-    async def test_send_message_returns_terminal_task(self):
-        """Test send_message when it receives a task in terminal state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.completed
-        
-        # Mock the async generator to yield a terminal task
-        async def mock_generator():
-            yield (mock_task,)
-        
-        self.mock_client.send_message.return_value = mock_generator()
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message(test_message)
-        
-        self.assertEqual(result, mock_task)
-
-    async def test_send_message_returns_interrupted_task(self):
-        """Test send_message when it receives a task in interrupted state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.input_required
-        
-        # Mock the async generator to yield an interrupted task
-        async def mock_generator():
-            yield (mock_task,)
-        
-        self.mock_client.send_message.return_value = mock_generator()
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message(test_message)
-        
-        self.assertEqual(result, mock_task)
-
-    async def test_send_message_returns_last_task(self):
-        """Test send_message returns the last task when no terminal/interrupted state is reached."""
-        mock_task1 = MagicMock(spec=Task)
-        mock_task1.status = MagicMock(spec=TaskStatus)
-        mock_task1.status.state = TaskState.working
-        
-        mock_task2 = MagicMock(spec=Task)
-        mock_task2.status = MagicMock(spec=TaskStatus)
-        mock_task2.status.state = TaskState.working
-        
-        # Mock the async generator to yield multiple running tasks
-        async def mock_generator():
-            yield (mock_task1,)
-            yield (mock_task2,)
-        
-        self.mock_client.send_message.return_value = mock_generator()
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message(test_message)
-        
-        self.assertEqual(result, mock_task2)
-
-    async def test_send_message_handles_exception(self):
-        """Test send_message properly handles and re-raises exceptions."""
-        test_error = Exception("Test error")
-        self.mock_client.send_message.side_effect = test_error
-        
-        test_message = MagicMock(spec=Message)
-        
-        with self.assertRaises(Exception) as context:
-            await self.connection.send_message(test_message)
-        
-        self.assertEqual(context.exception, test_error)
-
-    async def test_send_message_returns_none_when_no_events(self):
-        """Test send_message returns None when no events are yielded."""
-        # Mock empty async generator
-        async def mock_generator():
-            return
-            yield  # This line will never be reached
-        
-        self.mock_client.send_message.return_value = mock_generator()
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message(test_message)
-        
-        self.assertIsNone(result)
-
-    async def test_send_message_with_polling_returns_message(self):
-        """Test send_message_with_polling when send_message returns a Message."""
-        mock_message = MagicMock(spec=Message)
-        
-        # Mock send_message to return a Message
-        self.connection.send_message = AsyncMock(return_value=mock_message)
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message_with_polling(test_message)
-        
-        self.assertEqual(result, mock_message)
-        self.connection.send_message.assert_called_once_with(test_message)
-
-    async def test_send_message_with_polling_returns_terminal_task(self):
-        """Test send_message_with_polling when send_message returns a task in terminal state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.completed
-        
-        # Mock send_message to return a terminal task
-        self.connection.send_message = AsyncMock(return_value=mock_task)
-        
-        test_message = MagicMock(spec=Message)
-        result = await self.connection.send_message_with_polling(test_message)
-        
-        self.assertEqual(result, mock_task)
-        self.connection.send_message.assert_called_once_with(test_message)
-
-    async def test_send_message_with_polling_polls_until_terminal(self):
-        """Test send_message_with_polling polls until task reaches terminal state."""
-        # Create initial task in working state
-        initial_task = MagicMock(spec=Task)
-        initial_task.id = "task123"
-        initial_task.status = MagicMock(spec=TaskStatus)
-        initial_task.status.state = TaskState.working
-        
-        # Create intermediate task still in working state
-        intermediate_task = MagicMock(spec=Task)
-        intermediate_task.status = MagicMock(spec=TaskStatus)
-        intermediate_task.status.state = TaskState.working
-        
-        # Create final task in completed state
-        final_task = MagicMock(spec=Task)
-        final_task.status = MagicMock(spec=TaskStatus)
-        final_task.status.state = TaskState.completed
-        
-        # Mock send_message to return initial working task
-        self.connection.send_message = AsyncMock(return_value=initial_task)
-        
-        # Mock get_task to return intermediate task first, then final task
-        self.mock_client.get_task.side_effect = [intermediate_task, final_task]
-        
-        test_message = MagicMock(spec=Message)
-        
-        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
-            result = await self.connection.send_message_with_polling(
-                test_message, sleep_time=0.1, max_iter=10
-            )
-        
-        self.assertEqual(result, final_task)
-        self.connection.send_message.assert_called_once_with(test_message)
-        
-        # Verify get_task was called with correct parameters
-        expected_calls = [
-            call(TaskQueryParams(id="task123")),
-            call(TaskQueryParams(id="task123"))
-        ]
-        self.mock_client.get_task.assert_has_calls(expected_calls)
-        
-        # Verify sleep was called twice (once for each polling iteration)
-        self.assertEqual(mock_sleep.call_count, 2)
-        mock_sleep.assert_called_with(0.1)
-
-    async def test_send_message_with_polling_polls_until_interrupted(self):
-        """Test send_message_with_polling polls until task reaches interrupted state."""
-        # Create initial task in working state
-        initial_task = MagicMock(spec=Task)
-        initial_task.id = "task456"
-        initial_task.status = MagicMock(spec=TaskStatus)
-        initial_task.status.state = TaskState.working
-        
-        # Create final task in input_required state (interrupted)
-        final_task = MagicMock(spec=Task)
-        final_task.status = MagicMock(spec=TaskStatus)
-        final_task.status.state = TaskState.input_required
-        
-        # Mock send_message to return initial working task
-        self.connection.send_message = AsyncMock(return_value=initial_task)
-        
-        # Mock get_task to return final task
-        self.mock_client.get_task.return_value = final_task
-        
-        test_message = MagicMock(spec=Message)
-        
-        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
-            result = await self.connection.send_message_with_polling(
-                test_message, sleep_time=0.05, max_iter=5
-            )
-        
-        self.assertEqual(result, final_task)
-        self.mock_client.get_task.assert_called_once_with(TaskQueryParams(id="task456"))
-        mock_sleep.assert_called_once_with(0.05)
-
-    async def test_send_message_with_polling_timeout_exception(self):
-        """Test send_message_with_polling raises exception when max_iter is reached."""
-        # Create initial task in working state
-        initial_task = MagicMock(spec=Task)
-        initial_task.id = "task789"
-        initial_task.status = MagicMock(spec=TaskStatus)
-        initial_task.status.state = TaskState.working
-        
-        # Create task that stays in working state
-        working_task = MagicMock(spec=Task)
-        working_task.status = MagicMock(spec=TaskStatus)
-        working_task.status.state = TaskState.working
-        
-        # Mock send_message to return initial working task
-        self.connection.send_message = AsyncMock(return_value=initial_task)
-        
-        # Mock get_task to always return working task
-        self.mock_client.get_task.return_value = working_task
-        
-        test_message = MagicMock(spec=Message)
-        
-        with patch('asyncio.sleep', new_callable=AsyncMock):
-            with self.assertRaises(Exception) as context:
-                await self.connection.send_message_with_polling(
-                    test_message, sleep_time=0.1, max_iter=3
-                )
-        
-        expected_timeout = 3 * 0.1  # max_iter * sleep_time
-        self.assertIn(f"Task did not complete in {expected_timeout} seconds", str(context.exception))
-        
-        # Verify get_task was called max_iter times
-        self.assertEqual(self.mock_client.get_task.call_count, 3)
-
-    async def test_send_message_with_polling_no_task_returned(self):
-        """Test send_message_with_polling raises ValueError when send_message returns None."""
-        # Mock send_message to return None
-        self.connection.send_message = AsyncMock(return_value=None)
-        
-        test_message = MagicMock(spec=Message)
-        
-        with self.assertRaises(ValueError) as context:
-            await self.connection.send_message_with_polling(test_message)
-        
-        self.assertEqual(str(context.exception), "No task or message returned from send_message")
-        self.connection.send_message.assert_called_once_with(test_message)
-
-    async def test_send_message_with_polling_custom_parameters(self):
-        """Test send_message_with_polling with custom sleep_time and max_iter."""
-        # Create initial task in working state
-        initial_task = MagicMock(spec=Task)
-        initial_task.id = "custom_task"
-        initial_task.status = MagicMock(spec=TaskStatus)
-        initial_task.status.state = TaskState.working
-        
-        # Create final task in completed state
-        final_task = MagicMock(spec=Task)
-        final_task.status = MagicMock(spec=TaskStatus)
-        final_task.status.state = TaskState.completed
-        
-        # Mock send_message to return initial working task
-        self.connection.send_message = AsyncMock(return_value=initial_task)
-        
-        # Mock get_task to return final task
-        self.mock_client.get_task.return_value = final_task
-        
-        test_message = MagicMock(spec=Message)
-        
-        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
-            result = await self.connection.send_message_with_polling(
-                test_message, sleep_time=0.5, max_iter=100
-            )
-        
-        self.assertEqual(result, final_task)
-        mock_sleep.assert_called_once_with(0.5)
-
-    async def test_send_message_with_polling_default_parameters(self):
-        """Test send_message_with_polling uses default parameters correctly."""
-        # Create initial task in working state
-        initial_task = MagicMock(spec=Task)
-        initial_task.id = "default_task"
-        initial_task.status = MagicMock(spec=TaskStatus)
-        initial_task.status.state = TaskState.working
-        
-        # Create final task in failed state (terminal)
-        final_task = MagicMock(spec=Task)
-        final_task.status = MagicMock(spec=TaskStatus)
-        final_task.status.state = TaskState.failed
-        
-        # Mock send_message to return initial working task
-        self.connection.send_message = AsyncMock(return_value=initial_task)
-        
-        # Mock get_task to return final task
-        self.mock_client.get_task.return_value = final_task
-        
-        test_message = MagicMock(spec=Message)
-        
-        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
-            result = await self.connection.send_message_with_polling(test_message)
-        
-        self.assertEqual(result, final_task)
-        # Default sleep_time is 0.2
-        mock_sleep.assert_called_once_with(0.2)
-
-
-class TestHelperFunctions(unittest.TestCase):
-    """Tests for helper functions in the remote agent connection module."""
-
-    def test_is_in_terminal_state_completed(self):
-        """Test is_in_terminal_state returns True for completed state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.completed
-        
-        result = is_in_terminal_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_state_canceled(self):
-        """Test is_in_terminal_state returns True for canceled state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.canceled
-        
-        result = is_in_terminal_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_state_failed(self):
-        """Test is_in_terminal_state returns True for failed state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.failed
-        
-        result = is_in_terminal_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_state_running(self):
-        """Test is_in_terminal_state returns False for running state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.working
-        
-        result = is_in_terminal_state(mock_task)
-        self.assertFalse(result)
-
-    def test_is_in_terminal_or_interrupted_state_input_required(self):
-        """Test is_in_terminal_or_interrupted_state returns True for input_required state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.input_required
-        
-        result = is_in_terminal_or_interrupted_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_or_interrupted_state_unknown(self):
-        """Test is_in_terminal_or_interrupted_state returns True for unknown state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.unknown
-        
-        result = is_in_terminal_or_interrupted_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_or_interrupted_state_completed(self):
-        """Test is_in_terminal_or_interrupted_state returns True for terminal states."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.completed
-        
-        result = is_in_terminal_or_interrupted_state(mock_task)
-        self.assertTrue(result)
-
-    def test_is_in_terminal_or_interrupted_state_running(self):
-        """Test is_in_terminal_or_interrupted_state returns False for running state."""
-        mock_task = MagicMock(spec=Task)
-        mock_task.status = MagicMock(spec=TaskStatus)
-        mock_task.status.state = TaskState.working
-        
-        result = is_in_terminal_or_interrupted_state(mock_task)
-        self.assertFalse(result)
-
-
-if __name__ == '__main__':
-    unittest.main()