diff --git a/cognee/tests/unit/shared/test_usage_logger.py b/cognee/tests/unit/shared/test_usage_logger.py new file mode 100644 index 000000000..54b49454c --- /dev/null +++ b/cognee/tests/unit/shared/test_usage_logger.py @@ -0,0 +1,216 @@ +"""Unit tests for usage logger core functions.""" + +import pytest +from datetime import datetime, timezone +from uuid import UUID +from types import SimpleNamespace + +from cognee.shared.usage_logger import ( + _sanitize_value, + _sanitize_dict_key, + _get_param_names, + _get_param_defaults, + _extract_user_id, + _extract_parameters, + log_usage, +) +from cognee.shared.exceptions import UsageLoggerError + + +class TestSanitizeValue: + """Test _sanitize_value function.""" + + @pytest.mark.parametrize("value,expected", [ + (None, None), + ("string", "string"), + (42, 42), + (3.14, 3.14), + (True, True), + (False, False), + ]) + def test_basic_types(self, value, expected): + assert _sanitize_value(value) == expected + + def test_uuid_and_datetime(self): + """Test UUID and datetime serialization.""" + uuid_val = UUID("123e4567-e89b-12d3-a456-426614174000") + dt = datetime(2024, 1, 15, 12, 30, 45, tzinfo=timezone.utc) + + assert _sanitize_value(uuid_val) == "123e4567-e89b-12d3-a456-426614174000" + assert _sanitize_value(dt) == "2024-01-15T12:30:45+00:00" + + def test_collections(self): + """Test list, tuple, and dict serialization.""" + assert _sanitize_value([1, "string", UUID("123e4567-e89b-12d3-a456-426614174000"), None]) == [ + 1, "string", "123e4567-e89b-12d3-a456-426614174000", None + ] + assert _sanitize_value((1, "string", True)) == [1, "string", True] + assert _sanitize_value({"key": UUID("123e4567-e89b-12d3-a456-426614174000")}) == { + "key": "123e4567-e89b-12d3-a456-426614174000" + } + assert _sanitize_value([]) == [] + assert _sanitize_value({}) == {} + + def test_nested_and_complex(self): + """Test nested structures and non-serializable types.""" + # Nested structure + nested = {"level1": {"level2": {"level3": [1, 2, {"nested": "value"}]}}} + assert _sanitize_value(nested)["level1"]["level2"]["level3"][2]["nested"] == "value" + + # Non-serializable + class CustomObject: + def __str__(self): + return "" + + result = _sanitize_value(CustomObject()) + assert isinstance(result, str) + assert "" + + result = _sanitize_dict_key(BadKey()) + assert isinstance(result, str) + assert "