diff --git a/src/jules_agent_sdk/async_client.py b/src/jules_agent_sdk/async_client.py index 7ba29ce..4ae3214 100644 --- a/src/jules_agent_sdk/async_client.py +++ b/src/jules_agent_sdk/async_client.py @@ -106,6 +106,21 @@ async def wait_for_completion( await asyncio.sleep(poll_interval) + async def list_all(self) -> List[Session]: + """List all sessions asynchronously (handles pagination).""" + all_sessions: List[Session] = [] + page_token: Optional[str] = None + + while True: + result = await self.list(page_token=page_token) + all_sessions.extend(result["sessions"]) + + page_token = result.get("nextPageToken") + if not page_token: + break + + return all_sessions + class AsyncActivitiesAPI: """Async API client for managing session activities.""" diff --git a/src/jules_agent_sdk/sessions.py b/src/jules_agent_sdk/sessions.py index 93e3b50..afc10cb 100644 --- a/src/jules_agent_sdk/sessions.py +++ b/src/jules_agent_sdk/sessions.py @@ -194,3 +194,27 @@ def wait_for_completion( raise TimeoutError(f"Session polling timed out after {timeout} seconds") time.sleep(poll_interval) + + def list_all(self) -> List[Session]: + """List all sessions (handles pagination). + + Returns: + A list of all Session objects + + Example: + >>> all_sessions = client.sessions.list_all() + >>> for session in all_sessions: + ... print(session.id) + """ + all_sessions: List[Session] = [] + page_token: Optional[str] = None + + while True: + result = self.list(page_token=page_token) + all_sessions.extend(result["sessions"]) + + page_token = result.get("nextPageToken") + if not page_token: + break + + return all_sessions diff --git a/tests/test_async_client.py b/tests/test_async_client.py index 7a7af5b..eee6937 100644 --- a/tests/test_async_client.py +++ b/tests/test_async_client.py @@ -69,6 +69,65 @@ async def test_async_sessions_list(self, mock_request): assert len(result["sessions"]) == 1 assert result["sessions"][0].id == "test1" + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sessions_get(self, mock_request): + """Test async getting a session.""" + mock_request.return_value = {"id": "s1"} + client = AsyncJulesClient(api_key="test-api-key") + session = await client.sessions.get("s1") + assert session.id == "s1" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sessions_approve_plan(self, mock_request): + """Test async approving a plan.""" + client = AsyncJulesClient(api_key="test-api-key") + await client.sessions.approve_plan("s1") + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("POST", "sessions/s1:approvePlan") + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sessions_send_message(self, mock_request): + """Test async sending a message.""" + client = AsyncJulesClient(api_key="test-api-key") + await client.sessions.send_message("s1", "Test") + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("POST", "sessions/s1:sendMessage") + assert mock_request.call_args[1]["json"] == {"prompt": "Test"} + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_client.asyncio.sleep", new_callable=AsyncMock) + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sessions_wait_for_completion(self, mock_request, mock_sleep): + """Test async waiting for completion.""" + mock_request.side_effect = [ + {"state": "IN_PROGRESS"}, + {"state": "COMPLETED", "id": "s1"}, + ] + client = AsyncJulesClient(api_key="test-api-key") + session = await client.sessions.wait_for_completion("s1") + assert session.id == "s1" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_activities_get(self, mock_request): + """Test async getting an activity.""" + mock_request.return_value = {"id": "a1"} + client = AsyncJulesClient(api_key="test-api-key") + activity = await client.activities.get("s1", "a1") + assert activity.id == "a1" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_activities_list(self, mock_request): + """Test async listing activities.""" + mock_request.return_value = {"activities": [{"id": "a1"}]} + client = AsyncJulesClient(api_key="test-api-key") + result = await client.activities.list("s1") + assert len(result["activities"]) == 1 + @pytest.mark.asyncio @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") async def test_async_activities_list_all(self, mock_request): @@ -93,3 +152,72 @@ async def test_async_activities_list_all(self, mock_request): assert len(activities) == 2 assert activities[0].id == "a1" assert activities[1].id == "a2" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sources_get(self, mock_request): + """Test async getting a source.""" + mock_request.return_value = {"id": "src1"} + client = AsyncJulesClient(api_key="test-api-key") + source = await client.sources.get("src1") + assert source.id == "src1" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sources_list(self, mock_request): + """Test async listing sources.""" + mock_request.return_value = {"sources": [{"id": "src1"}]} + client = AsyncJulesClient(api_key="test-api-key") + result = await client.sources.list() + assert len(result["sources"]) == 1 + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sessions_list_all(self, mock_request): + """Test async listing all sessions with pagination.""" + mock_request.side_effect = [ + { + "sessions": [{"id": "s1"}], + "nextPageToken": "next", + }, + {"sessions": [{"id": "s2"}]}, + ] + client = AsyncJulesClient(api_key="test-api-key") + sessions = await client.sessions.list_all() + assert len(sessions) == 2 + assert sessions[0].id == "s1" + assert sessions[1].id == "s2" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.AsyncBaseClient._request") + async def test_async_sources_list_all(self, mock_request): + """Test async listing all sources with pagination.""" + mock_request.side_effect = [ + {"sources": [{"id": "src1"}], "nextPageToken": "next"}, + {"sources": [{"id": "src2"}]}, + ] + client = AsyncJulesClient(api_key="test-api-key") + sources = await client.sources.list_all() + assert len(sources) == 2 + assert sources[0].id == "src1" + assert sources[1].id == "src2" + + +class TestAsyncErrorHandling: + """Test error handling for the async client.""" + + @pytest.mark.asyncio + @patch("jules_agent_sdk.async_base.aiohttp.ClientSession.request") + async def test_async_authentication_error(self, mock_request): + """Test async authentication error.""" + mock_response = AsyncMock() + mock_response.ok = False + mock_response.status = 401 + mock_response.json.return_value = {"error": {"message": "Invalid API key"}} + mock_request.return_value.__aenter__.return_value = mock_response + + client = AsyncJulesClient(api_key="invalid-key") + with pytest.raises(JulesAuthenticationError): + await client.sessions.list() + + await client.close() diff --git a/tests/test_client.py b/tests/test_client.py index f36c07b..2a3512e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -115,6 +115,104 @@ def test_activities_list(self, mock_request): assert len(result["activities"]) == 2 assert result["activities"][0].id == "a1" + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sessions_approve_plan(self, mock_request): + """Test approving a session plan.""" + client = JulesClient(api_key="test-api-key") + client.sessions.approve_plan("s1") + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("POST", "sessions/s1:approvePlan") + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sessions_send_message(self, mock_request): + """Test sending a message to a session.""" + client = JulesClient(api_key="test-api-key") + client.sessions.send_message("s1", "Hello") + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("POST", "sessions/s1:sendMessage") + assert mock_request.call_args[1]["json"] == {"prompt": "Hello"} + + @patch("jules_agent_sdk.sessions.time.sleep", return_value=None) + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sessions_wait_for_completion_success(self, mock_request, mock_sleep): + """Test waiting for session completion successfully.""" + mock_request.side_effect = [ + {"state": "IN_PROGRESS"}, + {"state": "COMPLETED", "id": "s1"}, + ] + client = JulesClient(api_key="test-api-key") + session = client.sessions.wait_for_completion("s1") + assert session.id == "s1" + assert mock_request.call_count == 2 + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sessions_list_all(self, mock_request): + """Test listing all sessions with pagination.""" + mock_request.side_effect = [ + { + "sessions": [{"id": "s1"}], + "nextPageToken": "next", + }, + {"sessions": [{"id": "s2"}]}, + ] + client = JulesClient(api_key="test-api-key") + sessions = client.sessions.list_all() + assert len(sessions) == 2 + assert sessions[0].id == "s1" + assert sessions[1].id == "s2" + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_activities_get(self, mock_request): + """Test getting a single activity.""" + mock_request.return_value = {"id": "a1", "description": "Activity 1"} + client = JulesClient(api_key="test-api-key") + activity = client.activities.get("s1", "a1") + assert activity.id == "a1" + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("GET", "sessions/s1/activities/a1") + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_activities_list_all(self, mock_request): + """Test listing all activities with pagination.""" + mock_request.side_effect = [ + { + "activities": [{"id": "a1"}], + "nextPageToken": "next", + }, + {"activities": [{"id": "a2"}]}, + ] + client = JulesClient(api_key="test-api-key") + activities = client.activities.list_all("s1") + assert len(activities) == 2 + assert activities[0].id == "a1" + assert activities[1].id == "a2" + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sources_get(self, mock_request): + """Test getting a single source.""" + mock_request.return_value = {"id": "src1", "githubRepo": {"owner": "test"}} + client = JulesClient(api_key="test-api-key") + source = client.sources.get("src1") + assert source.id == "src1" + mock_request.assert_called_once() + assert mock_request.call_args[0] == ("GET", "sources/src1") + + @patch("jules_agent_sdk.base.BaseClient._request") + def test_sources_list_all(self, mock_request): + """Test listing all sources with pagination.""" + mock_request.side_effect = [ + { + "sources": [{"id": "src1"}], + "nextPageToken": "next", + }, + {"sources": [{"id": "src2"}]}, + ] + client = JulesClient(api_key="test-api-key") + sources = client.sources.list_all() + assert len(sources) == 2 + assert sources[0].id == "src1" + assert sources[1].id == "src2" + @patch("jules_agent_sdk.base.BaseClient._request") def test_sources_list(self, mock_request): """Test listing sources."""