diff --git a/quotientai/resources/logs.py b/quotientai/resources/logs.py index 6542113..2bf6b17 100644 --- a/quotientai/resources/logs.py +++ b/quotientai/resources/logs.py @@ -8,7 +8,7 @@ from collections import deque from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timezone from threading import Thread, Event from typing import Any, Dict, List, Optional, Union @@ -200,7 +200,7 @@ def create( log_id = str(uuid.uuid4()) # Create current timestamp - created_at = datetime.now().isoformat() + created_at = datetime.now(timezone.utc).isoformat() data = { "id": log_id, @@ -493,7 +493,7 @@ async def create( log_id = str(uuid.uuid4()) # Create current timestamp - created_at = datetime.now().isoformat() + created_at = datetime.now(timezone.utc).isoformat() # Create a copy of all the data data = { diff --git a/tests/test_client.py b/tests/test_client.py index 7aebaa7..4324d93 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -8,6 +8,7 @@ from quotientai.client import QuotientAI, QuotientLogger, _BaseQuotientClient + # Modify existing fixtures to use proper paths @pytest.fixture def mock_token_dir(tmp_path): @@ -16,10 +17,12 @@ def mock_token_dir(tmp_path): token_dir.mkdir(parents=True) return token_dir + @pytest.fixture def mock_api_key(): return "test-api-key" + @pytest.fixture def mock_auth_response(): return {"id": "test-user-id", "email": "test@example.com"} @@ -27,8 +30,9 @@ def mock_auth_response(): @pytest.fixture def mock_client(mock_auth_response): - with patch('quotientai.client._BaseQuotientClient') as MockClient, \ - patch('quotientai.client.AuthResource') as MockAuthResource: + with patch("quotientai.client._BaseQuotientClient") as MockClient, patch( + "quotientai.resources.AuthResource" + ) as MockAuthResource: mock_instance = MockClient.return_value mock_auth = MockAuthResource.return_value mock_auth.authenticate = Mock() @@ -37,69 +41,77 @@ def mock_client(mock_auth_response): mock_instance._delete = Mock() yield mock_instance + @pytest.fixture def quotient_client(mock_api_key, mock_client): - with patch.dict('os.environ', {'QUOTIENT_API_KEY': mock_api_key}): + with patch.dict("os.environ", {"QUOTIENT_API_KEY": mock_api_key}): client = QuotientAI() return client + class TestBaseQuotientClient: """Tests for the _BaseQuotientClient class""" - + def test_initialization(self, tmp_path): """Test the base client sets up correctly""" api_key = "test-api-key" - + # Use a clean temporary directory for token storage token_dir = tmp_path / ".quotient" - - with patch('pathlib.Path.home', return_value=tmp_path): + + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient(api_key) - + assert client.api_key == api_key assert client.token is None assert client.token_expiry == 0 assert client.token_api_key is None assert client.headers["Authorization"] == f"Bearer {api_key}" - assert client._token_path == tmp_path / ".quotient" / f"{api_key[-6:]}_auth_token.json" + assert ( + client._token_path + == tmp_path / ".quotient" / f"{api_key[-6:]}_auth_token.json" + ) def test_handle_jwt_response(self): """Test that _handle_response properly processes JWT tokens""" test_token = "test.jwt.token" test_expiry = int(time.time()) + 3600 - - with patch('jwt.decode') as mock_decode, \ - patch.object(_BaseQuotientClient, '_save_token') as mock_save_token: - + + with patch("jwt.decode") as mock_decode, patch.object( + _BaseQuotientClient, "_save_token" + ) as mock_save_token: + mock_decode.return_value = {"exp": test_expiry} - + client = _BaseQuotientClient("test-api-key") response = Mock() response.headers = {"X-JWT-Token": test_token} - + client._handle_response(response) - + # Verify _save_token was called with correct parameters mock_save_token.assert_called_once_with(test_token, test_expiry) - + # Verify the headers were updated assert client.headers["Authorization"] == f"Bearer {test_token}" def test_save_token(self, tmp_path): """Test that _save_token writes token data correctly""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") test_token = "test.jwt.token" test_expiry = int(time.time()) + 3600 - + client._save_token(test_token, test_expiry) - + # Verify token was saved in memory assert client.token == test_token assert client.token_expiry == test_expiry - + # Verify token was saved to disk - token_file = tmp_path / ".quotient" / f"{client.api_key[-6:]}_auth_token.json" + token_file = ( + tmp_path / ".quotient" / f"{client.api_key[-6:]}_auth_token.json" + ) assert token_file.exists() stored_data = json.loads(token_file.read_text()) assert stored_data["token"] == test_token @@ -108,24 +120,28 @@ def test_save_token(self, tmp_path): def test_load_token(self, tmp_path): """Test that _load_token reads token data correctly""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") test_token = "test.jwt.token" test_expiry = int(time.time()) + 3600 - + # Write a token file token_dir = tmp_path / ".quotient" token_dir.mkdir(parents=True) token_file = token_dir / f"{client.api_key[-6:]}_auth_token.json" - token_file.write_text(json.dumps({ - "token": test_token, - "expires_at": test_expiry, - "api_key": client.api_key - })) - + token_file.write_text( + json.dumps( + { + "token": test_token, + "expires_at": test_expiry, + "api_key": client.api_key, + } + ) + ) + # Load the token client._load_token() - + assert client.token == test_token assert client.token_expiry == test_expiry assert client.token_api_key == client.api_key @@ -133,30 +149,30 @@ def test_load_token(self, tmp_path): def test_is_token_valid(self, tmp_path): """Test token validity checking""" # Prevent token loading by using clean temp directory - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - + # Test with no token assert not client._is_token_valid() - + # Test with expired token client.token = "expired.token" client.token_expiry = int(time.time()) - 3600 # 1 hour ago client.token_api_key = client.api_key assert not client._is_token_valid() - + # Test with valid token client.token = "valid.token" client.token_expiry = int(time.time()) + 3600 # 1 hour from now client.token_api_key = client.api_key assert client._is_token_valid() - + # Test with token about to expire (within 5 minute buffer) client.token = "about.to.expire" client.token_expiry = int(time.time()) + 200 # Less than 5 minutes client.token_api_key = client.api_key assert not client._is_token_valid() - + # Test with mismatched API key client.token = "valid.token" client.token_expiry = int(time.time()) + 3600 @@ -166,13 +182,13 @@ def test_is_token_valid(self, tmp_path): def test_update_auth_header(self, tmp_path): """Test authorization header updates based on token state""" # Prevent token loading by using clean temp directory - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - + # Should use API key when no token client._update_auth_header() assert client.headers["Authorization"] == f"Bearer {client.api_key}" - + # Should use valid token test_token = "test.jwt.token" client.token = test_token @@ -180,12 +196,12 @@ def test_update_auth_header(self, tmp_path): client.token_api_key = client.api_key client._update_auth_header() assert client.headers["Authorization"] == f"Bearer {test_token}" - + # Should revert to API key when token expires client.token_expiry = int(time.time()) - 3600 client._update_auth_header() assert client.headers["Authorization"] == f"Bearer {client.api_key}" - + # Should revert to API key when API key doesn't match client.token = test_token client.token_expiry = int(time.time()) + 3600 @@ -194,65 +210,70 @@ def test_update_auth_header(self, tmp_path): assert client.headers["Authorization"] == f"Bearer {client.api_key}" def test_token_path_uses_home(self): - with patch('pathlib.Path.home') as mock_home: - mock_home.return_value = Path('/home/user') - client = _BaseQuotientClient('test-key') - assert client._token_path == Path('/home/user/.quotient/st-key_auth_token.json') + with patch("pathlib.Path.home") as mock_home: + mock_home.return_value = Path("/home/user") + client = _BaseQuotientClient("test-key") + assert client._token_path == Path( + "/home/user/.quotient/st-key_auth_token.json" + ) def test_token_path_fallback_to_root(self): - with patch('pathlib.Path.home') as mock_home, \ - patch('pathlib.Path.exists') as mock_exists: + with patch("pathlib.Path.home") as mock_home, patch( + "pathlib.Path.exists" + ) as mock_exists: # Simulate home directory access failing mock_home.side_effect = Exception("Can't access home") # Simulate /root existing mock_exists.return_value = True - - client = _BaseQuotientClient('test-key') - assert client._token_path == Path('/root/.quotient/st-key_auth_token.json') + + client = _BaseQuotientClient("test-key") + assert client._token_path == Path("/root/.quotient/st-key_auth_token.json") def test_token_path_fallback_to_cwd(self): - with patch('pathlib.Path.home') as mock_home, \ - patch('pathlib.Path.exists') as mock_exists, \ - patch('pathlib.Path.cwd') as mock_cwd: + with patch("pathlib.Path.home") as mock_home, patch( + "pathlib.Path.exists" + ) as mock_exists, patch("pathlib.Path.cwd") as mock_cwd: # Simulate home directory access failing mock_home.side_effect = Exception("Can't access home") # Simulate /root not existing mock_exists.return_value = False # Set current working directory - mock_cwd.return_value = Path('/current/dir') - - client = _BaseQuotientClient('test-key') - assert client._token_path == Path('/current/dir/.quotient/st-key_auth_token.json') + mock_cwd.return_value = Path("/current/dir") + + client = _BaseQuotientClient("test-key") + assert client._token_path == Path( + "/current/dir/.quotient/st-key_auth_token.json" + ) def test_handle_jwt_token_success(self): - client = _BaseQuotientClient('test-key') + client = _BaseQuotientClient("test-key") mock_response = Mock() test_token = "test.jwt.token" test_expiry = int(time.time()) + 3600 - + # Create a real JWT token with an expiry mock_response.headers = {"X-JWT-Token": test_token} - - with patch('jwt.decode') as mock_decode: + + with patch("jwt.decode") as mock_decode: mock_decode.return_value = {"exp": test_expiry} client._handle_response(mock_response) - + assert client.headers["Authorization"] == f"Bearer {test_token}" assert client.token == test_token assert client.token_expiry == test_expiry def test_handle_jwt_token_decode_error(self): - client = _BaseQuotientClient('test-key') + client = _BaseQuotientClient("test-key") mock_response = Mock() original_auth = client.headers["Authorization"] test_token = "test.jwt.token" test_expiry = int(time.time()) + 3600 mock_response.headers = {"X-JWT-Token": test_token} - - with patch('jwt.decode') as mock_decode: - mock_decode.side_effect = jwt.InvalidTokenError("Invalid token") + + with patch("jwt.decode") as mock_decode: + mock_decode.side_effect = Exception("Invalid token") client._handle_response(mock_response) - + # Should keep original authorization header on error assert client.headers["Authorization"] == original_auth # Token is still saved even if decoding fails @@ -261,8 +282,9 @@ def test_handle_jwt_token_decode_error(self): def test_token_dir_creation_error(self, tmp_path, caplog): """Test handling of directory creation errors""" - with patch('pathlib.Path.home', return_value=tmp_path), \ - patch('pathlib.Path.mkdir', side_effect=Exception("Permission denied")): + with patch("pathlib.Path.home", return_value=tmp_path), patch( + "pathlib.Path.mkdir", side_effect=Exception("Permission denied") + ): client = _BaseQuotientClient("test-api-key") result = client._save_token("test-token", int(time.time()) + 3600) assert result is None @@ -271,78 +293,87 @@ def test_token_dir_creation_error(self, tmp_path, caplog): def test_get_wrapper(self, tmp_path): """Test _get wrapper method""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - with patch.object(client, 'get') as mock_get: + with patch.object(client, "get") as mock_get: mock_response = Mock() mock_get.return_value = mock_response - + # Test with params client._get("/test", params={"key": "value"}, timeout=30) - mock_get.assert_called_with("/test", params={"key": "value"}, timeout=30) - + mock_get.assert_called_with( + "/test", params={"key": "value"}, timeout=30 + ) + # Test without params client._get("/test") mock_get.assert_called_with("/test", params=None, timeout=None) def test_post_wrapper(self, tmp_path): """Test _post wrapper method""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - with patch.object(client, 'post') as mock_post: + with patch.object(client, "post") as mock_post: mock_response = Mock() mock_post.return_value = mock_response - + # Test with data data = {"key": "value", "null_key": None} client._post("/test", data=data, timeout=30) - mock_post.assert_called_with(url="/test", json={"key": "value"}, timeout=30) - + mock_post.assert_called_with( + url="/test", json={"key": "value"}, timeout=30 + ) + # Test with list data list_data = ["value", None, "value2"] client._post("/test", data=list_data) - mock_post.assert_called_with(url="/test", json=["value", "value2"], timeout=None) - + mock_post.assert_called_with( + url="/test", json=["value", "value2"], timeout=None + ) + # Test without data client._post("/test") mock_post.assert_called_with(url="/test", json={}, timeout=None) def test_patch_wrapper(self, tmp_path): """Test _patch wrapper method""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - with patch.object(client, 'patch') as mock_patch: + with patch.object(client, "patch") as mock_patch: mock_response = Mock() mock_patch.return_value = mock_response - + # Test with data data = {"key": "value", "null_key": None} client._patch("/test", data=data, timeout=30) - mock_patch.assert_called_with(url="/test", json={"key": "value"}, timeout=30) - + mock_patch.assert_called_with( + url="/test", json={"key": "value"}, timeout=30 + ) + # Test without data client._patch("/test") mock_patch.assert_called_with(url="/test", json={}, timeout=None) def test_delete_wrapper(self, tmp_path): """Test _delete wrapper method""" - with patch('pathlib.Path.home', return_value=tmp_path): + with patch("pathlib.Path.home", return_value=tmp_path): client = _BaseQuotientClient("test-api-key") - with patch.object(client, 'delete') as mock_delete: + with patch.object(client, "delete") as mock_delete: mock_response = Mock() mock_delete.return_value = mock_response - + # Test with timeout client._delete("/test", timeout=30) mock_delete.assert_called_with("/test", timeout=30) - + # Test without timeout client._delete("/test") mock_delete.assert_called_with("/test", timeout=None) + class TestQuotientAI: """Tests for the QuotientAI class""" - + def test_init_with_api_key(self, mock_client): api_key = "test-api-key" client = QuotientAI(api_key=api_key) @@ -350,36 +381,42 @@ def test_init_with_api_key(self, mock_client): def test_init_with_env_var(self, mock_client): api_key = "test-api-key" - with patch.dict('os.environ', {'QUOTIENT_API_KEY': api_key}): + with patch.dict("os.environ", {"QUOTIENT_API_KEY": api_key}): client = QuotientAI() assert client.api_key == api_key def test_init_no_api_key(self, caplog): - with patch.dict('os.environ', clear=True): + with patch.dict("os.environ", clear=True): client = QuotientAI() assert client.api_key is None assert "could not find API key" in caplog.text assert "https://app.quotientai.co" in caplog.text def test_init_calls_auth_resource(self, mock_client): - with patch('quotientai.client.AuthResource') as MockAuthResource: + with patch("quotientai.resources.AuthResource") as MockAuthResource: api_key = "test-api-key" client = QuotientAI(api_key=api_key) assert MockAuthResource.call_count == 1 def test_init_auth_failure(self, caplog): """Test that the client logs authentication failure""" - with patch('quotientai.client.AuthResource') as MockAuthResource: + with patch("quotientai.resources.AuthResource") as MockAuthResource: mock_auth = MockAuthResource.return_value - mock_auth.authenticate = Mock(side_effect=Exception("'Exception' object has no attribute 'message'")) + mock_auth.authenticate = Mock( + side_effect=Exception("'Exception' object has no attribute 'message'") + ) QuotientAI(api_key="test-api-key") assert "'Exception' object has no attribute 'message'" in caplog.text - assert "If you are seeing this error, please check that your API key is correct" in caplog.text + assert ( + "If you are seeing this error, please check that your API key is correct" + in caplog.text + ) + class TestQuotientLogger: """Tests for the QuotientLogger class""" - + def test_init(self): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) @@ -389,15 +426,15 @@ def test_init(self): def test_configuration(self): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) - + logger.init( app_name="test-app", environment="test", tags={"tag1": "value1"}, sample_rate=0.5, - hallucination_detection=True + hallucination_detection=True, ) - + assert logger._configured assert logger.app_name == "test-app" assert logger.environment == "test" @@ -409,11 +446,7 @@ def test_invalid_sample_rate(self, caplog): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) - result = logger.init( - app_name="test-app", - environment="test", - sample_rate=2.0 - ) + result = logger.init(app_name="test-app", environment="test", sample_rate=2.0) assert result is None assert "sample_rate must be between 0.0 and 1.0" in caplog.text @@ -422,12 +455,9 @@ def test_invalid_sample_rate(self, caplog): def test_log_without_init(self, caplog): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) - - result = logger.log( - user_query="test query", - model_output="test output" - ) - + + result = logger.log(user_query="test query", model_output="test output") + assert result is None assert "Logger is not configured" in caplog.text @@ -435,8 +465,8 @@ def test_log_with_init(self): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) logger.init(app_name="test-app", environment="test") - - with patch.object(QuotientLogger, '_should_sample', return_value=True): + + with patch.object(QuotientLogger, "_should_sample", return_value=True): logger.log(user_query="test query", model_output="test output") assert mock_logs_resource.create.call_count == 1 @@ -444,8 +474,8 @@ def test_log_respects_sampling(self): mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) logger.init(app_name="test-app", environment="test") - - with patch.object(QuotientLogger, '_should_sample', return_value=False): + + with patch.object(QuotientLogger, "_should_sample", return_value=False): logger.log(user_query="test query", model_output="test output") assert mock_logs_resource.create.call_count == 0 @@ -453,29 +483,29 @@ def test_should_sample(self): """Test sampling logic with different sample rates""" mock_logs_resource = Mock() logger = QuotientLogger(mock_logs_resource) - - with patch('random.random') as mock_random: + + with patch("random.random") as mock_random: # Test with 100% sample rate logger.sample_rate = 1.0 mock_random.return_value = 0.5 assert logger._should_sample() is True - + # Test with 0% sample rate logger.sample_rate = 0.0 mock_random.return_value = 0.5 assert logger._should_sample() is False - + # Test with 50% sample rate logger.sample_rate = 0.5 - + # Should sample when random < sample_rate mock_random.return_value = 0.4 assert logger._should_sample() is True - + # Should not sample when random >= sample_rate mock_random.return_value = 0.6 assert logger._should_sample() is False - + def test_log_with_invalid_document_dict(self, caplog): """Test logging with an invalid document dictionary""" mock_logs_resource = Mock() @@ -485,7 +515,7 @@ def test_log_with_invalid_document_dict(self, caplog): result = logger.log( user_query="test query", model_output="test output", - documents=[{"metadata": {"key": "value"}}] + documents=[{"metadata": {"key": "value"}}], ) assert result is None @@ -493,7 +523,7 @@ def test_log_with_invalid_document_dict(self, caplog): assert "page_content" in caplog.text assert mock_logs_resource.create.call_count == 0 assert "Documents must include 'page_content' field" in caplog.text - + def test_log_with_invalid_document_type(self, caplog): """Test logging with a document of invalid type""" mock_logs_resource = Mock() @@ -503,7 +533,7 @@ def test_log_with_invalid_document_type(self, caplog): result = logger.log( user_query="test query", model_output="test output", - documents=["valid string document", 123] + documents=["valid string document", 123], ) assert result is None @@ -520,7 +550,7 @@ def test_log_with_valid_documents(self): logger.init(app_name="test-app", environment="test") # Force sampling to True for testing - with patch.object(logger, '_should_sample', return_value=True): + with patch.object(logger, "_should_sample", return_value=True): # Test with mixed valid document types logger.log( user_query="test query 4", @@ -528,7 +558,7 @@ def test_log_with_valid_documents(self): documents=[ "string document", {"page_content": "dict document", "metadata": {"key": "value"}}, - ] + ], ) assert mock_logs_resource.create.call_count == 1 @@ -536,5 +566,8 @@ def test_log_with_valid_documents(self): # Verify correct documents were passed to create calls = mock_logs_resource.create.call_args_list assert calls[0][1]["documents"][0] == "string document" - assert calls[0][1]["documents"][1] == {"page_content": "dict document", "metadata": {"key": "value"}} + assert calls[0][1]["documents"][1] == { + "page_content": "dict document", + "metadata": {"key": "value"}, + } assert len(calls[0][1]["documents"]) == 2