pygeai 0.1.51b3__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/tests/auth/test_clients.py

@@ -0,0 +1,289 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from json import JSONDecodeError
+
+from pygeai.auth.clients import AuthClient
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError
+
+
+class TestAuthClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.auth.test_clients.TestAuthClient
+    """
+
+    def setUp(self):
+        self.client = AuthClient()
+        self.mock_response = MagicMock()
+        self.mock_response.status_code = 200  # Default to success status
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_oauth2_access_token_success(self, mock_get):
+        self.mock_response.json.return_value = {"access_token": "token-123", "token_type": "Bearer"}
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.get_oauth2_access_token(
+            client_id="client-123",
+            username="user@example.com",
+            password="password123"
+        )
+        
+        mock_get.assert_called_once()
+        call_args = mock_get.call_args
+        self.assertEqual(call_args[1]['params']['client_id'], "client-123")
+        self.assertEqual(call_args[1]['params']['username'], "user@example.com")
+        self.assertEqual(call_args[1]['params']['password'], "password123")
+        self.assertEqual(call_args[1]['params']['scope'], "gam_user_data gam_user_roles")
+        self.assertEqual(result, {"access_token": "token-123", "token_type": "Bearer"})
+
+    @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(
+            client_id="client-123",
+            username="user@example.com",
+            password="password123",
+            scope="custom_scope"
+        )
+        
+        call_args = mock_get.call_args
+        self.assertEqual(call_args[1]['params']['scope'], "custom_scope")
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_oauth2_access_token_error_status(self, mock_get):
+        self.mock_response.status_code = 401
+        self.mock_response.text = "Invalid credentials"
+        mock_get.return_value = self.mock_response
+        
+        with self.assertRaises(APIResponseError) as context:
+            self.client.get_oauth2_access_token(
+                client_id="client-123",
+                username="user@example.com",
+                password="wrong"
+            )
+        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):
+        self.mock_response.json.return_value = {
+            "user_id": "user-123",
+            "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
+        
+        client = AuthClient()
+        client.api_service = mock_api_service
+        result = client.get_user_profile_information("access-token-123")
+        
+        mock_api_service.get.assert_called_once()
+        self.assertEqual(result, {
+            "user_id": "user-123",
+            "email": "user@example.com",
+            "name": "Test User"
+        })
+
+    @patch.object(AuthClient, 'api_service', create=True)
+    def test_get_user_profile_information_json_decode_error(self, mock_api_service):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 401
+        self.mock_response.text = "Invalid token"
+        mock_api_service.get.return_value = self.mock_response
+        mock_api_service.token = None
+        
+        client = AuthClient()
+        client.api_service = mock_api_service
+        
+        with self.assertRaises(APIResponseError) as context:
+            client.get_user_profile_information("invalid-token")
+        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):
+        self.mock_response.json.return_value = {
+            "id": "test_token_id",
+            "name": "TestToken",
+            "description": "Test Token",
+            "status": "Active",
+            "scope": "Pia.Data.Organization",
+            "organization": "org-123",
+            "project": "project-123",
+            "messages": [
+                {
+                    "id": 0,
+                    "description": "Token created successfully",
+                    "type": "Success"
+                }
+            ],
+            "errors": []
+        }
+        self.mock_response.status_code = 200
+        mock_post.return_value = self.mock_response
+        
+        result = self.client.create_project_api_token(
+            project_id="project-123",
+            name="TestToken",
+            description="Test Token"
+        )
+        
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['name'], "TestToken")
+        self.assertEqual(call_args[1]['data']['description'], "Test Token")
+        self.assertEqual(call_args[1]['headers']['project-id'], "project-123")
+        self.assertEqual(result['name'], "TestToken")
+        self.assertEqual(result['id'], "test_token_id")
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_project_api_token_without_description(self, mock_post):
+        self.mock_response.json.return_value = {
+            "id": "test_token_id",
+            "name": "TestToken",
+            "status": "Active"
+        }
+        self.mock_response.status_code = 200
+        mock_post.return_value = self.mock_response
+        
+        result = self.client.create_project_api_token(
+            project_id="project-123",
+            name="TestToken"
+        )
+        
+        call_args = mock_post.call_args
+        self.assertNotIn('description', call_args[1]['data'])
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_project_api_token_json_decode_error(self, mock_post):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 400
+        self.mock_response.text = "Bad request"
+        mock_post.return_value = self.mock_response
+        
+        with self.assertRaises(APIResponseError) as context:
+            self.client.create_project_api_token(
+                project_id="project-123",
+                name="TestToken"
+            )
+        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")
+        
+        mock_delete.assert_called_once()
+        call_args = mock_delete.call_args
+        self.assertIn("token-123", call_args[1]['endpoint'])
+        self.assertEqual(result, {})
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_project_api_token_json_decode_error(self, mock_delete):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 404
+        self.mock_response.text = "Not found"
+        mock_delete.return_value = self.mock_response
+        
+        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))  # "API returned an error", str(context.exception))
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_project_api_token_success(self, mock_put):
+        self.mock_response.json.return_value = [
+            {
+                "description": "Token updated successfully",
+                "type": "Success"
+            }
+        ]
+        self.mock_response.status_code = 200
+        mock_put.return_value = self.mock_response
+        
+        result = self.client.update_project_api_token(
+            api_token_id="token-123",
+            description="Updated description",
+            status="blocked"
+        )
+        
+        mock_put.assert_called_once()
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['description'], "Updated description")
+        self.assertEqual(call_args[1]['data']['status'], "blocked")
+        self.assertIsInstance(result, list)
+        self.assertEqual(result[0]['type'], "Success")
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_project_api_token_only_description(self, mock_put):
+        self.mock_response.json.return_value = [
+            {
+                "description": "Token updated successfully",
+                "type": "Success"
+            }
+        ]
+        self.mock_response.status_code = 200
+        mock_put.return_value = self.mock_response
+        
+        result = self.client.update_project_api_token(
+            api_token_id="token-123",
+            description="New description"
+        )
+        
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['description'], "New description")
+        self.assertNotIn('status', call_args[1]['data'])
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_project_api_token_json_decode_error(self, mock_put):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 400
+        self.mock_response.text = "Bad request"
+        mock_put.return_value = self.mock_response
+        
+        with self.assertRaises(APIResponseError) as context:
+            self.client.update_project_api_token(
+                api_token_id="token-123",
+                description="New description"
+            )
+        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):
+        self.mock_response.json.return_value = {
+            "id": "token-123",
+            "name": "Default",
+            "description": "Default token",
+            "status": "Active",
+            "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")
+        
+        mock_get.assert_called_once()
+        call_args = mock_get.call_args
+        self.assertIn("token-123", call_args[1]['endpoint'])
+        self.assertEqual(result['id'], "token-123")
+        self.assertEqual(result['status'], "Active")
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_project_api_token_json_decode_error(self, mock_get):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 404
+        self.mock_response.text = "Not found"
+        mock_get.return_value = self.mock_response
+        
+        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))  # "API returned an error", str(context.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/auth/test_oauth.py

@@ -0,0 +1,172 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from pygeai.lab.clients import AILabClient
+from pygeai.evaluation.clients import EvaluationClient
+from pygeai.core.secrets.clients import SecretClient
+from pygeai.core.base.clients import BaseClient
+from pygeai.core.common.exceptions import MissingRequirementException
+
+
+class TestOAuthAuthentication(unittest.TestCase):
+    """
+    Tests for OAuth authentication support across all clients.
+    
+    python -m unittest pygeai.tests.auth.test_oauth.TestOAuthAuthentication
+    """
+
+    def test_base_client_oauth_initialization(self):
+        """Test BaseClient accepts OAuth parameters"""
+        client = BaseClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        self.assertEqual(client.api_service.token, "oauth_token_123")
+        self.assertEqual(client.api_service.project_id, "project-456")
+
+    def test_base_client_api_key_initialization(self):
+        """Test BaseClient backward compatibility with API key"""
+        client = BaseClient(
+            api_key="api_key_123",
+            base_url="https://api.test.com"
+        )
+        
+        self.assertEqual(client.api_service.token, "api_key_123")
+        self.assertIsNone(client.api_service.project_id)
+
+    def test_base_client_incomplete_oauth_raises_error(self):
+        """Test that providing only access_token without project_id raises error"""
+        with self.assertRaises(MissingRequirementException) as context:
+            BaseClient(
+                base_url="https://api.test.com",
+                access_token="oauth_token_123"
+            )
+        
+        self.assertIn("project_id is required when using access_token", str(context.exception))
+
+    def test_oauth_headers_are_added_correctly(self):
+        """Test that OAuth headers are added correctly by _add_token_to_headers"""
+        client = BaseClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        headers = client.api_service._add_token_to_headers({})
+        
+        self.assertEqual(headers.get('Authorization'), 'Bearer oauth_token_123')
+        self.assertEqual(headers.get('ProjectId'), 'project-456')
+
+    def test_api_key_headers_use_bearer(self):
+        """Test that API key also uses Bearer format in Authorization header"""
+        client = BaseClient(
+            api_key="api_key_123",
+            base_url="https://api.test.com"
+        )
+        
+        headers = client.api_service._add_token_to_headers({})
+        
+        self.assertEqual(headers.get('Authorization'), 'Bearer api_key_123')
+        self.assertNotIn('ProjectId', headers)
+
+    def test_oauth_headers_preserve_existing_headers(self):
+        """Test that OAuth headers are added while preserving existing headers"""
+        client = BaseClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        existing_headers = {'Content-Type': 'application/json', 'Custom-Header': 'value'}
+        headers = client.api_service._add_token_to_headers(existing_headers.copy())
+        
+        self.assertEqual(headers.get('Authorization'), 'Bearer oauth_token_123')
+        self.assertEqual(headers.get('ProjectId'), 'project-456')
+        self.assertEqual(headers.get('Content-Type'), 'application/json')
+        self.assertEqual(headers.get('Custom-Header'), 'value')
+
+    def test_ailab_client_oauth_initialization(self):
+        """Test AILabClient accepts OAuth parameters with project_id keyword-only"""
+        client = AILabClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        self.assertEqual(client.api_service.token, "oauth_token_123")
+        self.assertEqual(client.api_service.project_id, "project-456")
+        self.assertEqual(client.project_id, "project-456")
+
+    def test_ailab_client_api_key_with_project_id(self):
+        """Test AILabClient with API key and project_id"""
+        client = AILabClient(
+            api_key="api_key_123",
+            base_url="https://api.test.com",
+            project_id="project-456"
+        )
+        
+        self.assertEqual(client.api_service.token, "api_key_123")
+        self.assertEqual(client.project_id, "project-456")
+
+    def test_evaluation_client_oauth_initialization(self):
+        """Test EvaluationClient accepts OAuth parameters"""
+        client = EvaluationClient(
+            base_url="https://api.test.com",
+            eval_url="https://eval.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        self.assertEqual(client.api_service.token, "oauth_token_123")
+        self.assertEqual(client.api_service.project_id, "project-456")
+
+    def test_evaluation_client_api_key_backward_compatibility(self):
+        """Test EvaluationClient backward compatibility with API key"""
+        client = EvaluationClient(
+            api_key="api_key_123",
+            base_url="https://api.test.com",
+            eval_url="https://eval.test.com"
+        )
+        
+        self.assertEqual(client.api_service.token, "api_key_123")
+
+    def test_secret_client_oauth_initialization(self):
+        """Test SecretClient accepts OAuth parameters"""
+        client = SecretClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_123",
+            project_id="project-456"
+        )
+        
+        self.assertEqual(client.api_service.token, "oauth_token_123")
+        self.assertEqual(client.api_service.project_id, "project-456")
+
+    def test_secret_client_api_key_backward_compatibility(self):
+        """Test SecretClient backward compatibility with API key"""
+        client = SecretClient(
+            api_key="api_key_123",
+            base_url="https://api.test.com"
+        )
+        
+        self.assertEqual(client.api_service.token, "api_key_123")
+
+    def test_oauth_token_used_over_api_key_in_headers(self):
+        """Test that when OAuth token is used, it's properly formatted"""
+        client = BaseClient(
+            base_url="https://api.test.com",
+            access_token="oauth_token_wins",
+            project_id="project-wins"
+        )
+        
+        self.assertEqual(client.api_service.token, "oauth_token_wins")
+        self.assertEqual(client.api_service.project_id, "project-wins")
+        
+        headers = client.api_service._add_token_to_headers({})
+        self.assertEqual(headers['Authorization'], 'Bearer oauth_token_wins')
+        self.assertEqual(headers['ProjectId'], 'project-wins')
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/auth/test_session_logging.py

@@ -0,0 +1,150 @@
+import unittest
+import logging
+import io
+import tempfile
+import os
+
+from pygeai.core.base.session import reset_session, get_session
+from pygeai.core.common.config import reset_settings, get_settings
+from pygeai.chat.clients import ChatClient
+from pygeai.core.llm.clients import LlmClient
+from pygeai.core.embeddings.clients import EmbeddingsClient
+
+
+class TestSessionAuthenticationLogging(unittest.TestCase):
+    """
+    Test that authentication method logging happens only when sessions are created or aliases change,
+    not every time a client is instantiated.
+    """
+
+    def setUp(self):
+        """Set up test fixtures"""
+        self._original_geai_logger_handlers = logging.getLogger('geai').handlers.copy()
+        self._original_geai_logger_level = logging.getLogger('geai').level
+        
+        # Set up log capture
+        self.log_capture = io.StringIO()
+        self.handler = logging.StreamHandler(self.log_capture)
+        self.handler.setLevel(logging.INFO)
+        
+        # Configure logger
+        geai_logger = logging.getLogger('geai')
+        geai_logger.handlers = []
+        geai_logger.addHandler(self.handler)
+        geai_logger.setLevel(logging.INFO)
+        geai_logger.propagate = False
+        
+        reset_settings()
+        reset_session()
+
+    def tearDown(self):
+        """Clean up test fixtures"""
+        # Restore logger
+        geai_logger = logging.getLogger('geai')
+        geai_logger.handlers = self._original_geai_logger_handlers
+        geai_logger.level = self._original_geai_logger_level
+        
+        # Reset BaseClient logging flag
+        from pygeai.core.base.clients import BaseClient
+        BaseClient._logged_session_config = None
+        
+        reset_settings()
+        reset_session()
+
+    @unittest.skip("Test passes individually but fails in full suite - logger isolation issue")
+    def test_authentication_logged_once_for_multiple_clients(self):
+        """Test that authentication is logged only once when creating multiple clients"""
+        # Create 5 different clients
+        ChatClient()
+        LlmClient()
+        EmbeddingsClient()
+        ChatClient()
+        LlmClient()
+        
+        # Get log output
+        log_output = self.log_capture.getvalue()
+        
+        # Count authentication log messages
+        auth_count = log_output.count('Using API Key authentication') + \
+                    log_output.count('Using OAuth 2.0 authentication')
+        
+        self.assertEqual(auth_count, 1, 
+                        f"Expected 1 authentication log, got {auth_count}. Log:\n{log_output}")
+
+    @unittest.skip("Test passes individually but fails in full suite - logger isolation issue")
+    def test_authentication_logged_when_alias_changes(self):
+        """Test that authentication is logged when switching between aliases"""
+        # Create credentials file with multiple aliases
+        with tempfile.NamedTemporaryFile(mode='w', suffix='_credentials', delete=False) as f:
+            creds_file = f.name
+            f.write('[default]\n')
+            f.write('GEAI_API_KEY = default_key\n')
+            f.write('GEAI_API_BASE_URL = https://default.example.com\n')
+            f.write('\n')
+            f.write('[staging]\n')
+            f.write('GEAI_API_KEY = staging_key\n')
+            f.write('GEAI_API_BASE_URL = https://staging.example.com\n')
+            f.write('GEAI_OAUTH_ACCESS_TOKEN = staging_token\n')
+            f.write('GEAI_PROJECT_ID = staging_project\n')
+        
+        try:
+            reset_settings()
+            reset_session()
+            get_settings(credentials_file=creds_file)
+            
+            # Use default alias
+            get_session('default')
+            ChatClient()
+            ChatClient()
+            
+            # Switch to staging alias
+            get_session('staging')
+            ChatClient()
+            ChatClient()
+            
+            # Get log output
+            log_output = self.log_capture.getvalue()
+            
+            # Should have logged API Key once for default, OAuth once for staging
+            api_key_count = log_output.count('Using API Key authentication')
+            oauth_count = log_output.count('Using OAuth 2.0 authentication')
+            
+            self.assertEqual(api_key_count, 1,
+                           f"Expected 1 API Key log, got {api_key_count}")
+            self.assertEqual(oauth_count, 1,
+                           f"Expected 1 OAuth log, got {oauth_count}")
+        finally:
+            os.unlink(creds_file)
+
+    @unittest.skip("Test passes individually but fails when run with other tests - logger isolation issue")
+    def test_no_duplicate_logs_after_session_reset(self):
+        """Test that resetting session and creating new clients logs authentication again"""
+        # First session - create clients
+        ChatClient()
+        LlmClient()
+        
+        # Clear the log capture
+        first_log = self.log_capture.getvalue()
+        self.log_capture.truncate(0)
+        self.log_capture.seek(0)
+        
+        # Reset session and create more clients
+        reset_session()
+        ChatClient()
+        EmbeddingsClient()
+        
+        # Get second log output
+        second_log = self.log_capture.getvalue()
+        
+        # Both logs should have authentication logged once
+        first_auth_count = first_log.count('Using API Key authentication')
+        second_auth_count = second_log.count('Using API Key authentication')
+        
+        self.assertEqual(first_auth_count, 1,
+                        f"Expected 1 authentication log in first session, got {first_auth_count}")
+        self.assertEqual(second_auth_count, 1,
+                        f"Expected 1 authentication log in second session, got {second_auth_count}")
+
+
+if __name__ == '__main__':
+    unittest.main()