c3d33e728e
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin --------- Co-authored-by: Daniel Molnar <soobrosa@gmail.com> Co-authored-by: Igor Ilic <30923996+dexters1@users.noreply.github.com> Co-authored-by: Igor Ilic <igorilic03@gmail.com>
122 lines
4 KiB
Python
122 lines
4 KiB
Python
import unittest
|
|
import os
|
|
import uuid
|
|
from unittest.mock import patch, MagicMock
|
|
import sys
|
|
|
|
# Import the telemetry function to test
|
|
from cognee.shared.utils import send_telemetry
|
|
|
|
|
|
class TestTelemetry(unittest.TestCase):
|
|
@patch("cognee.shared.utils.requests.post")
|
|
def test_telemetry_enabled(self, mock_post):
|
|
# Setup mock response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_post.return_value = mock_response
|
|
|
|
# Check if .anon_id exists in the project root
|
|
anon_id_path = os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(__file__))), ".anon_id"
|
|
)
|
|
|
|
if not os.path.exists(anon_id_path):
|
|
# Create the file with a test ID if it doesn't exist
|
|
with open(anon_id_path, "w") as f:
|
|
f.write("test-machine")
|
|
print(f"Created .anon_id file at {anon_id_path}", file=sys.stderr)
|
|
|
|
self.assertTrue(os.path.exists(anon_id_path), "The .anon_id file should exist")
|
|
|
|
# Verify the file has content
|
|
with open(anon_id_path, "r") as f:
|
|
content = f.read().strip()
|
|
|
|
self.assertTrue(len(content) > 0, "The .anon_id file should not be empty")
|
|
|
|
# Ensure telemetry is enabled for this test
|
|
if "TELEMETRY_DISABLED" in os.environ:
|
|
del os.environ["TELEMETRY_DISABLED"]
|
|
|
|
# Make sure ENV is not test or dev
|
|
original_env = os.environ.get("ENV")
|
|
os.environ["ENV"] = "prod" # Set to dev to ensure telemetry is sent
|
|
|
|
# Generate a random user ID for testing
|
|
test_user_id = str(uuid.uuid4())
|
|
|
|
# Test sending telemetry
|
|
event_name = "test_event"
|
|
additional_props = {"test_key": "test_value"}
|
|
|
|
send_telemetry(event_name, test_user_id, additional_props)
|
|
|
|
# Verify telemetry was sent
|
|
mock_post.assert_called_once()
|
|
|
|
# Get the args that were passed to post
|
|
args, kwargs = mock_post.call_args
|
|
|
|
# Check that the payload contains our data
|
|
self.assertIn("json", kwargs)
|
|
payload = kwargs["json"]
|
|
|
|
# Verify payload contains expected data
|
|
self.assertEqual(payload.get("event_name"), event_name)
|
|
|
|
# Check that user_id is in the correct nested structure
|
|
self.assertIn("user_properties", payload)
|
|
self.assertEqual(payload["user_properties"].get("user_id"), str(test_user_id))
|
|
|
|
# Also check that user_id is in the properties
|
|
self.assertIn("properties", payload)
|
|
self.assertEqual(payload["properties"].get("user_id"), str(test_user_id))
|
|
|
|
# Check that additional properties are included
|
|
self.assertEqual(payload["properties"].get("test_key"), "test_value")
|
|
|
|
# Restore original ENV if it existed
|
|
if original_env is not None:
|
|
os.environ["ENV"] = original_env
|
|
else:
|
|
del os.environ["ENV"]
|
|
|
|
@patch("cognee.shared.utils.requests.post")
|
|
def test_telemetry_disabled(self, mock_post):
|
|
# Enable the TELEMETRY_DISABLED environment variable
|
|
os.environ["TELEMETRY_DISABLED"] = "1"
|
|
|
|
# Test sending telemetry
|
|
send_telemetry("disabled_test", "user123", {"key": "value"})
|
|
|
|
# Verify telemetry was not sent
|
|
mock_post.assert_not_called()
|
|
|
|
# Clean up
|
|
del os.environ["TELEMETRY_DISABLED"]
|
|
|
|
@patch("cognee.shared.utils.requests.post")
|
|
def test_telemetry_dev_env(self, mock_post):
|
|
# Set ENV to dev which should disable telemetry
|
|
original_env = os.environ.get("ENV")
|
|
os.environ["ENV"] = "dev"
|
|
|
|
if "TELEMETRY_DISABLED" in os.environ:
|
|
del os.environ["TELEMETRY_DISABLED"]
|
|
|
|
# Test sending telemetry
|
|
send_telemetry("dev_test", "user123", {"key": "value"})
|
|
|
|
# Verify telemetry was not sent in dev environment
|
|
mock_post.assert_not_called()
|
|
|
|
# Restore original ENV if it existed
|
|
if original_env is not None:
|
|
os.environ["ENV"] = original_env
|
|
else:
|
|
del os.environ["ENV"]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|