pygeai 0.6.0b7__py3-none-any.whl → 0.6.0b10__py3-none-any.whl

pygeai/tests/assistants/rag/test_clients.py

@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 from pygeai.assistant.rag.clients import RAGAssistantClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
 
 
 class TestRAGAssistantClient(unittest.TestCase):
@@ -29,6 +29,7 @@ class TestRAGAssistantClient(unittest.TestCase):
     def test_get_assistants_from_project_successful_json_response(self):
         expected_response = {"assistants": []}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.get_assistants_from_project()
@@ -42,15 +43,16 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_assistants_from_project()
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to get assistants from project", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get assistants from project", str(context.exception))
 
     def test_get_assistant_data_successful_json_response(self):
         expected_response = {"name": self.assistant_name, "details": "test"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.get_assistant_data(self.assistant_name)
@@ -64,17 +66,18 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_assistant_data(self.assistant_name)
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to get assistant data", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get assistant data", str(context.exception))
 
     def test_create_assistant_successful_json_response(self):
         name = "new_assistant"
         description = "Test description"
         expected_response = {"id": "new_id", "name": name}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post.return_value = self.mock_response
 
         result = self.client.create_assistant(name, description)
@@ -90,17 +93,18 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.post.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.create_assistant(name, description)
 
         self.client.api_service.post.assert_called_once()
-        self.assertIn("Unable to create assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to create assistant", str(context.exception))
 
     def test_update_assistant_successful_json_response(self):
         name = "test_assistant"
         status = 1
         expected_response = {"name": name, "status": status}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.put.return_value = self.mock_response
 
         result = self.client.update_assistant(name, status)
@@ -116,16 +120,17 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.put.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.update_assistant(name, status)
 
         self.client.api_service.put.assert_called_once()
-        self.assertIn("Unable to update assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to update assistant", str(context.exception))
 
     def test_delete_assistant_successful_json_response(self):
         name = "test_assistant"
         expected_response = {"message": "Assistant deleted"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.delete.return_value = self.mock_response
 
         result = self.client.delete_assistant(name)
@@ -140,11 +145,11 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.delete.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.delete_assistant(name)
 
         self.client.api_service.delete.assert_called_once()
-        self.assertIn("Unable to delete assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to delete assistant", str(context.exception))
 
     def test_get_documents_successful_json_response(self):
         name = "test_assistant"
@@ -152,6 +157,7 @@ class TestRAGAssistantClient(unittest.TestCase):
         count = 10
         expected_response = {"documents": []}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.get_documents(name, skip, count)
@@ -168,16 +174,17 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_documents(name, skip, count)
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to get documents", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get documents", str(context.exception))
 
     def test_delete_all_documents_successful_json_response(self):
         name = "test_assistant"
         expected_response = {"message": "All documents deleted"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.delete.return_value = self.mock_response
 
         result = self.client.delete_all_documents(name)
@@ -192,17 +199,18 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.delete.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.delete_all_documents(name)
 
         self.client.api_service.delete.assert_called_once()
-        self.assertIn("Unable to delete all documents", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to delete all documents", str(context.exception))
 
     def test_retrieve_document_successful_json_response(self):
         name = "test_assistant"
         document_id = "doc_123"
         expected_response = {"id": document_id, "content": "test content"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.retrieve_document(name, document_id)
@@ -218,11 +226,11 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.retrieve_document(name, document_id)
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to retrieve document", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve document", str(context.exception))
 
     def test_upload_document_binary_successful_json_response(self):
         name = "test_assistant"
@@ -231,6 +239,7 @@ class TestRAGAssistantClient(unittest.TestCase):
         content_type = "application/pdf"
         expected_response = {"id": "uploaded_doc", "status": "success"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post_file_binary.return_value = self.mock_response
 
         with patch('builtins.open', return_value=MagicMock()):
@@ -247,6 +256,7 @@ class TestRAGAssistantClient(unittest.TestCase):
         metadata = {"key": "value"}
         expected_response = {"id": "uploaded_doc", "status": "success"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post_files_multipart.return_value = self.mock_response
 
         with patch('builtins.open', return_value=MagicMock()), \
@@ -277,17 +287,18 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.client.api_service.post_file_binary.return_value = self.mock_response
 
         with patch('builtins.open', return_value=MagicMock()):
-            with self.assertRaises(InvalidAPIResponseException) as context:
+            with self.assertRaises(APIResponseError) as context:
                 self.client.upload_document(name, file_path, upload_type, content_type=content_type)
 
         self.client.api_service.post_file_binary.assert_called_once()
-        self.assertIn("Unable to upload document", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to upload document", str(context.exception))
 
     def test_delete_document_successful_json_response(self):
         name = "test_assistant"
         document_id = "doc_123"
         expected_response = {"message": "Document deleted"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.delete.return_value = self.mock_response
 
         result = self.client.delete_document(name, document_id)
@@ -303,16 +314,17 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.delete.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.delete_document(name, document_id)
 
         self.client.api_service.delete.assert_called_once()
-        self.assertIn("Unable to delete document", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to delete document", str(context.exception))
 
     def test_execute_query_successful_json_response(self):
         query = {"text": "test query"}
         expected_response = {"result": "query response"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post.return_value = self.mock_response
 
         result = self.client.execute_query(query)
@@ -327,8 +339,8 @@ class TestRAGAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.post.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.execute_query(query)
 
         self.client.api_service.post.assert_called_once()
-        self.assertIn("Unable to execute query", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to execute query", str(context.exception))

pygeai/tests/assistants/test_clients.py

@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 from pygeai.assistant.clients import AssistantClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
 
 
 class TestAssistantClient(unittest.TestCase):
@@ -22,6 +22,7 @@ class TestAssistantClient(unittest.TestCase):
         detail = "summary"
         expected_response = {"id": self.assistant_id, "name": "Test Assistant"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.get_assistant_data(self.assistant_id, detail)
@@ -36,11 +37,11 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_assistant_data(self.assistant_id, detail)
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to get assistant data", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get assistant data", str(context.exception))
 
     def test_create_assistant_successful_json_response(self):
         assistant_type = "text"
@@ -48,6 +49,7 @@ class TestAssistantClient(unittest.TestCase):
         prompt = "Help with text"
         expected_response = {"id": "new_id", "name": name}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post.return_value = self.mock_response
 
         result = self.client.create_assistant(assistant_type, name, prompt)
@@ -64,17 +66,18 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.post.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.create_assistant(assistant_type, name, prompt)
 
         self.client.api_service.post.assert_called_once()
-        self.assertIn("Unable to create assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to create assistant", str(context.exception))
 
     def test_update_assistant_successful_json_response(self):
         status = 1
         action = "save"
         expected_response = {"id": self.assistant_id, "status": status}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.put.return_value = self.mock_response
 
         result = self.client.update_assistant(self.assistant_id, status, action)
@@ -90,15 +93,16 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.put.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.update_assistant(self.assistant_id, status, action)
 
         self.client.api_service.put.assert_called_once()
-        self.assertIn("Unable to update assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to update assistant", str(context.exception))
 
     def test_delete_assistant_successful_json_response(self):
         expected_response = {"message": "Assistant deleted"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.delete.return_value = self.mock_response
 
         result = self.client.delete_assistant(self.assistant_id)
@@ -112,16 +116,17 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.delete.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.delete_assistant(self.assistant_id)
 
         self.client.api_service.delete.assert_called_once()
-        self.assertIn("Unable to delete assistant", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to delete assistant", str(context.exception))
 
     def test_send_chat_request_successful_json_response(self):
         messages = [{"role": "user", "content": "Hello"}]
         expected_response = {"response": "Hi there!"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post.return_value = self.mock_response
 
         result = self.client.send_chat_request(self.assistant_name, messages)
@@ -136,15 +141,16 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.post.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.send_chat_request(self.assistant_name, messages)
 
         self.client.api_service.post.assert_called_once()
-        self.assertIn("Unable to send chat request", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to send chat request", str(context.exception))
 
     def test_get_request_status_successful_json_response(self):
         expected_response = {"status": "completed"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.get.return_value = self.mock_response
 
         result = self.client.get_request_status(self.request_id)
@@ -158,15 +164,16 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.get.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_request_status(self.request_id)
 
         self.client.api_service.get.assert_called_once()
-        self.assertIn("Unable to get request status", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get request status", str(context.exception))
 
     def test_cancel_request_successful_json_response(self):
         expected_response = {"message": "Request cancelled"}
         self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
         self.client.api_service.post.return_value = self.mock_response
 
         result = self.client.cancel_request(self.request_id)
@@ -180,9 +187,9 @@ class TestAssistantClient(unittest.TestCase):
         self.mock_response.status_code = 500
         self.client.api_service.post.return_value = self.mock_response
 
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.cancel_request(self.request_id)
 
         self.client.api_service.post.assert_called_once()
-        self.assertIn("Unable to cancel request", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to cancel request", str(context.exception))
 

pygeai/tests/auth/test_clients.py

@@ -38,6 +38,7 @@ class TestAuthClient(unittest.TestCase):
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_oauth2_access_token_custom_scope(self, mock_get):
         self.mock_response.json.return_value = {"access_token": "token-123"}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_oauth2_access_token(
@@ -62,7 +63,7 @@ class TestAuthClient(unittest.TestCase):
                 username="user@example.com",
                 password="wrong"
             )
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
     @patch.object(AuthClient, 'api_service', create=True)
     def test_get_user_profile_information_success(self, mock_api_service):
@@ -71,6 +72,7 @@ class TestAuthClient(unittest.TestCase):
             "email": "user@example.com",
             "name": "Test User"
         }
+        self.mock_response.status_code = 200
         mock_api_service.get.return_value = self.mock_response
         mock_api_service.token = None
         
@@ -98,7 +100,7 @@ class TestAuthClient(unittest.TestCase):
         
         with self.assertRaises(APIResponseError) as context:
             client.get_user_profile_information("invalid-token")
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.post')
     def test_create_project_api_token_success(self, mock_post):
@@ -119,6 +121,7 @@ class TestAuthClient(unittest.TestCase):
             ],
             "errors": []
         }
+        self.mock_response.status_code = 200
         mock_post.return_value = self.mock_response
         
         result = self.client.create_project_api_token(
@@ -142,6 +145,7 @@ class TestAuthClient(unittest.TestCase):
             "name": "TestToken",
             "status": "Active"
         }
+        self.mock_response.status_code = 200
         mock_post.return_value = self.mock_response
         
         result = self.client.create_project_api_token(
@@ -164,11 +168,12 @@ class TestAuthClient(unittest.TestCase):
                 project_id="project-123",
                 name="TestToken"
             )
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.delete')
     def test_delete_project_api_token_success(self, mock_delete):
         self.mock_response.json.return_value = {}
+        self.mock_response.status_code = 200
         mock_delete.return_value = self.mock_response
         
         result = self.client.delete_project_api_token(api_token_id="token-123")
@@ -187,7 +192,7 @@ class TestAuthClient(unittest.TestCase):
         
         with self.assertRaises(APIResponseError) as context:
             self.client.delete_project_api_token(api_token_id="invalid-token")
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.put')
     def test_update_project_api_token_success(self, mock_put):
@@ -197,6 +202,7 @@ class TestAuthClient(unittest.TestCase):
                 "type": "Success"
             }
         ]
+        self.mock_response.status_code = 200
         mock_put.return_value = self.mock_response
         
         result = self.client.update_project_api_token(
@@ -220,6 +226,7 @@ class TestAuthClient(unittest.TestCase):
                 "type": "Success"
             }
         ]
+        self.mock_response.status_code = 200
         mock_put.return_value = self.mock_response
         
         result = self.client.update_project_api_token(
@@ -243,7 +250,7 @@ class TestAuthClient(unittest.TestCase):
                 api_token_id="token-123",
                 description="New description"
             )
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_project_api_token_success(self, mock_get):
@@ -255,6 +262,7 @@ class TestAuthClient(unittest.TestCase):
             "scope": "Pia.Data.Organization",
             "timestamp": "2024-07-22T18:37:32.341"
         }
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_project_api_token(api_token_id="token-123")
@@ -274,7 +282,7 @@ class TestAuthClient(unittest.TestCase):
         
         with self.assertRaises(APIResponseError) as context:
             self.client.get_project_api_token(api_token_id="invalid-token")
-        self.assertIn("API returned an error", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "API returned an error", str(context.exception))
 
 
 if __name__ == '__main__':