test cases for telemetry

Author: nithinjamespallicka

tests/telemetry/__init__.py

@@ -0,0 +1,42 @@
+import pytest
+
+from datu.telemetry.product.events import ProductTelemetryEvent
+
+
+@pytest.fixture
+def sample_event_data():
+    return {
+        "event_name": "test_event",
+        "user_id": "12345",
+        "timestamp": "2024-06-01T12:00:00Z",
+        "properties": {"plan": "pro", "source": "web"},
+    }
+
+
+@pytest.fixture
+def event(sample_event_data):
+    return ProductTelemetryEvent(
+        event_name=sample_event_data["event_name"],
+        user_id=sample_event_data["user_id"],
+        timestamp=sample_event_data["timestamp"],
+        properties=sample_event_data["properties"],
+    )
+
+
+@pytest.fixture
+def sample_event():
+    return ProductTelemetryEvent(foo="bar")
+
+
+@pytest.fixture
+def telemetry_settings():
+    from datu.telemetry.config import TelemetryConfig
+
+    return TelemetryConfig(api_key="dummy_key", package_name="datu-core")
+
+
+@pytest.fixture
+def posthog_client(telemetry_settings):
+    from datu.telemetry.product.posthog import PostHogClient
+
+    return PostHogClient(settings=telemetry_settings)

tests/telemetry/conftest.py

@@ -0,0 +1,75 @@
+import pytest
+
+from datu.telemetry.product.events import MCPClientEvent, OpenAIEvent, ProductTelemetryEvent
+
+
+def test_event_initialization(event, sample_event_data):
+    assert event.properties["event_name"] == sample_event_data["event_name"]
+    assert event.properties["user_id"] == sample_event_data["user_id"]
+    assert event.properties["timestamp"] == sample_event_data["timestamp"]
+    assert event.properties["properties"] == sample_event_data["properties"]
+
+
+def test_event_name_and_batch_key(event):
+    assert event.name == "ProductTelemetryEvent"
+    assert event.batch_key == event.name
+
+
+def test_batching_same_type(event):
+    other = ProductTelemetryEvent(event_name="other_event")
+    batched = event.batch(other)
+
+    # Batch size increments
+    assert batched.batch_size == 2
+    assert batched is event  # batching modifies self
+
+
+def test_batching_different_type_raises():
+    class AnotherEvent(ProductTelemetryEvent):
+        pass
+
+    e1 = ProductTelemetryEvent()
+    e2 = AnotherEvent()
+    with pytest.raises(ValueError):
+        e1.batch(e2)
+
+
+def test_batch_size_increment(event):
+    # Initial batch_size
+    assert event.batch_size == 1
+    event.batch(ProductTelemetryEvent())
+    assert event.batch_size == 2
+    event.batch(ProductTelemetryEvent())
+    assert event.batch_size == 3
+
+
+def test_mcp_client_event_properties():
+    servers = ["playwright", "puppeteer"]
+    event = MCPClientEvent(server_names=servers)
+
+    # Check properties
+    assert event.properties["mcp_server_names"] == servers
+    # Name and batch_key should come from base class
+    assert event.name == "MCPClientEvent"
+    assert event.batch_key == event.name
+    assert event.batch_size == 1
+
+
+def test_openai_event_properties():
+    from datu.app_config import get_app_settings
+
+    app_settings = get_app_settings()
+    data = {"user_id": "123", "action": "test"}
+    event = OpenAIEvent(**data)
+
+    # Base properties
+    for k, v in data.items():
+        assert event.properties[k] == v
+
+    # Extra property added in subclass
+    assert event.properties["openai_model"] == app_settings.openai_model
+
+    # Name and batch_key
+    assert event.name == "OpenAIEvent"
+    assert event.batch_key == event.name
+    assert event.batch_size == 1

tests/telemetry/product/__init__.py

@@ -0,0 +1,103 @@
+from unittest.mock import patch
+
+
+def test_posthog_client_initialization(posthog_client, telemetry_settings):
+    from datu.telemetry.product.posthog import PostHogClient
+
+    assert posthog_client.settings == telemetry_settings
+    assert isinstance(posthog_client._batched_events, dict)
+    assert isinstance(posthog_client.session_id, str)
+    assert posthog_client._user_id == ""
+    assert posthog_client._user_id_path == PostHogClient.USER_ID_PATH
+
+
+def test_user_id_creation(tmp_path):
+    from datu.telemetry.config import TelemetryConfig
+    from datu.telemetry.product.posthog import PostHogClient
+
+    path = tmp_path / "telemetry_user_id"
+    client = PostHogClient(settings=TelemetryConfig())
+    client._user_id_path = path
+
+    # file does not exist yet
+    uid = client.user_id
+    assert uid != PostHogClient.UNKNOWN_USER_ID
+    assert path.read_text().strip() == uid
+
+    # file exists, reads the same
+    uid2 = client.user_id
+    assert uid2 == uid
+
+
+def test_user_id_fallback_patch():
+    from pathlib import Path
+
+    from datu.telemetry.config import TelemetryConfig
+    from datu.telemetry.product.posthog import PostHogClient
+
+    client = PostHogClient(settings=TelemetryConfig())
+
+    with (
+        patch.object(Path, "exists", side_effect=OSError("fail")),
+        patch.object(Path, "read_text", side_effect=OSError("fail")),
+    ):
+        uid = client.user_id
+
+    assert uid == PostHogClient.UNKNOWN_USER_ID
+
+
+def test_base_context(monkeypatch):
+    from datu.telemetry.config import TelemetryConfig
+    from datu.telemetry.product.posthog import PostHogClient
+
+    client = PostHogClient(settings=TelemetryConfig(package_name="nonexistent_pkg"))
+
+    context = client._base_context()
+    assert "python_version" in context
+    assert "os" in context
+    assert "os_version" in context
+    assert context["package_version"] == "unknown"
+    assert isinstance(context["extras_installed"], dict)
+
+
+def test_capture_single_event(monkeypatch, posthog_client, sample_event):
+    # Ensure _send is called
+    called = {}
+
+    def fake_send(event):
+        called["event_name"] = event.name
+
+    posthog_client._send = fake_send
+
+    posthog_client.capture(sample_event)
+    assert called["event_name"] == sample_event.name
+
+
+def test_capture_batching(monkeypatch):
+    from datu.telemetry.config import TelemetryConfig
+    from datu.telemetry.product.events import ProductTelemetryEvent
+    from datu.telemetry.product.posthog import PostHogClient
+
+    settings = TelemetryConfig(api_key="dummy")
+    client = PostHogClient(settings=settings)
+
+    class BatchEvent(ProductTelemetryEvent):
+        max_batch_size = 2
+
+    e1 = BatchEvent(foo=1)
+    e2 = BatchEvent(foo=2)
+
+    sent = []
+
+    def fake_send(event):
+        sent.append(event)
+
+    client._send = fake_send
+
+    client.capture(e1)
+    assert client._batched_events[e1.batch_key].batch_size == 1
+    assert sent == []
+
+    client.capture(e2)
+    assert sent[0].batch_size == 2
+    assert e1.batch_key not in client._batched_events

tests/telemetry/product/conftest.py

@@ -0,0 +1,6 @@
+def test_config():
+    from datu.telemetry.config import get_telemetry_settings
+
+    settings = get_telemetry_settings()
+    assert settings.api_key == "phc_m74dfR9nLpm2nipvkL2swyFDtNuQNC9o2FL2CSbh6Je"
+    assert settings.package_name == "datu-core"