350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""
|
|
Unit tests for telemetry module.
|
|
|
|
Tests cover:
|
|
1. OFF state: no recording, no disk write
|
|
2. ON state: events buffered and saved
|
|
3. Buffer limit: oldest events purged
|
|
4. TTL: old events auto-purged
|
|
5. Invalid category: rejected
|
|
6. Invalid label: rejected
|
|
7. Unknown pattern ID: mapped to __unknown__
|
|
8. PII detection: rejected
|
|
9. Rate limiting: excess events rejected
|
|
10. Clear: buffer emptied
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from datetime import datetime, timedelta
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Add parent directory to path for imports
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from telemetry import (
|
|
TelemetryStore,
|
|
TelemetryEvent,
|
|
RateLimiter,
|
|
validate_event,
|
|
contains_pii,
|
|
get_pattern_allowlist,
|
|
validate_pattern_label,
|
|
SCHEMA_VERSION,
|
|
MAX_EVENTS,
|
|
MAX_LABEL_LENGTH,
|
|
EVENT_TTL_DAYS,
|
|
)
|
|
|
|
|
|
class TestRateLimiter(unittest.TestCase):
|
|
"""Test rate limiter functionality."""
|
|
|
|
def test_allows_initial_requests(self):
|
|
"""Rate limiter should allow initial burst."""
|
|
limiter = RateLimiter(max_per_minute=10)
|
|
for _ in range(10):
|
|
self.assertTrue(limiter.allow())
|
|
|
|
def test_rejects_after_limit(self):
|
|
"""Rate limiter should reject after limit reached."""
|
|
limiter = RateLimiter(max_per_minute=5)
|
|
for _ in range(5):
|
|
limiter.allow()
|
|
self.assertFalse(limiter.allow())
|
|
|
|
def test_refills_over_time(self):
|
|
"""Rate limiter should refill tokens over time."""
|
|
limiter = RateLimiter(max_per_minute=60)
|
|
# Exhaust all tokens
|
|
for _ in range(60):
|
|
limiter.allow()
|
|
self.assertFalse(limiter.allow())
|
|
|
|
# Simulate time passing (1 second = 1 token)
|
|
limiter.last_refill -= 2 # 2 seconds ago
|
|
self.assertTrue(limiter.allow())
|
|
|
|
|
|
class TestValidation(unittest.TestCase):
|
|
"""Test event validation."""
|
|
|
|
def test_valid_event(self):
|
|
"""Valid event should pass validation."""
|
|
is_valid, msg, event = validate_event({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
self.assertTrue(is_valid)
|
|
self.assertIsNotNone(event)
|
|
|
|
def test_missing_category(self):
|
|
"""Missing category should fail validation."""
|
|
is_valid, msg, _ = validate_event({
|
|
"action": "tab_switch",
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("category", msg.lower())
|
|
|
|
def test_invalid_category(self):
|
|
"""Invalid category should fail validation."""
|
|
is_valid, msg, _ = validate_event({
|
|
"category": "invalid_category",
|
|
"action": "tab_switch",
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("category", msg.lower())
|
|
|
|
def test_invalid_action(self):
|
|
"""Invalid action should fail validation."""
|
|
is_valid, msg, _ = validate_event({
|
|
"category": "feature",
|
|
"action": "invalid_action",
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("action", msg.lower())
|
|
|
|
def test_invalid_label(self):
|
|
"""Invalid label should fail validation."""
|
|
is_valid, msg, _ = validate_event({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "invalid_label",
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("label", msg.lower())
|
|
|
|
def test_label_too_long(self):
|
|
"""Label exceeding max length should fail."""
|
|
long_label = "a" * (MAX_LABEL_LENGTH + 1)
|
|
is_valid, msg, _ = validate_event({
|
|
"category": "session",
|
|
"action": "start",
|
|
"label": long_label,
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("long", msg.lower())
|
|
|
|
def test_value_out_of_range(self):
|
|
"""Value out of range should fail validation."""
|
|
is_valid, msg, _ = validate_event({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
"value": -1,
|
|
})
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("range", msg.lower())
|
|
|
|
|
|
class TestPIIDetection(unittest.TestCase):
|
|
"""Test PII pattern detection."""
|
|
|
|
def test_detects_email(self):
|
|
"""Should detect email addresses."""
|
|
self.assertTrue(contains_pii("user@example.com"))
|
|
|
|
def test_detects_windows_path(self):
|
|
"""Should detect Windows user paths."""
|
|
self.assertTrue(contains_pii(r"C:\Users\john\file.py"))
|
|
|
|
def test_detects_linux_home(self):
|
|
"""Should detect Linux home paths."""
|
|
self.assertTrue(contains_pii("/home/alice/test.py"))
|
|
|
|
def test_detects_macos_home(self):
|
|
"""Should detect macOS home paths."""
|
|
self.assertTrue(contains_pii("/Users/bob/documents"))
|
|
|
|
def test_detects_api_key(self):
|
|
"""Should detect API key patterns."""
|
|
self.assertTrue(contains_pii("sk-abc123def456ghi789jkl012mno345"))
|
|
|
|
def test_allows_safe_strings(self):
|
|
"""Should allow safe strings."""
|
|
self.assertFalse(contains_pii("chat"))
|
|
self.assertFalse(contains_pii("resolved"))
|
|
self.assertFalse(contains_pii("cuda_oom"))
|
|
|
|
|
|
class TestPatternValidation(unittest.TestCase):
|
|
"""Test pattern ID validation."""
|
|
|
|
@patch('telemetry.get_pattern_allowlist')
|
|
def test_valid_pattern_id(self, mock_allowlist):
|
|
"""Valid pattern ID should pass."""
|
|
mock_allowlist.return_value = {"cuda_oom", "missing_module"}
|
|
result = validate_pattern_label("cuda_oom")
|
|
self.assertEqual(result, "cuda_oom")
|
|
|
|
@patch('telemetry.get_pattern_allowlist')
|
|
def test_unknown_pattern_id(self, mock_allowlist):
|
|
"""Unknown pattern ID should map to __unknown__."""
|
|
mock_allowlist.return_value = {"cuda_oom"}
|
|
result = validate_pattern_label("unknown_pattern")
|
|
self.assertEqual(result, "__unknown__")
|
|
|
|
def test_empty_label(self):
|
|
"""Empty label should map to __unknown__."""
|
|
result = validate_pattern_label("")
|
|
self.assertEqual(result, "__unknown__")
|
|
|
|
|
|
class TestTelemetryStore(unittest.TestCase):
|
|
"""Test telemetry store functionality."""
|
|
|
|
def setUp(self):
|
|
"""Create temp file for each test."""
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.temp_file = os.path.join(self.temp_dir, "telemetry.json")
|
|
|
|
def tearDown(self):
|
|
"""Clean up temp files."""
|
|
if os.path.exists(self.temp_file):
|
|
os.remove(self.temp_file)
|
|
if os.path.exists(self.temp_dir):
|
|
os.rmdir(self.temp_dir)
|
|
|
|
def test_disabled_no_recording(self):
|
|
"""When disabled, no events should be recorded."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=False)
|
|
success, msg = store.track({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
self.assertFalse(success)
|
|
self.assertIn("disabled", msg.lower())
|
|
self.assertFalse(os.path.exists(self.temp_file))
|
|
|
|
def test_enabled_records_event(self):
|
|
"""When enabled, valid events should be recorded."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
success, msg = store.track({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
self.assertTrue(success)
|
|
self.assertTrue(os.path.exists(self.temp_file))
|
|
self.assertEqual(len(store), 1)
|
|
|
|
def test_buffer_limit(self):
|
|
"""Buffer should not exceed MAX_EVENTS."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
# Bypass rate limiter for this test
|
|
store._rate_limiter = RateLimiter(max_per_minute=10000)
|
|
|
|
# Add more than MAX_EVENTS
|
|
for i in range(MAX_EVENTS + 10):
|
|
store.track({
|
|
"category": "session",
|
|
"action": "start",
|
|
})
|
|
|
|
self.assertEqual(len(store), MAX_EVENTS)
|
|
|
|
def test_clear(self):
|
|
"""Clear should empty the buffer."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
store.track({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
self.assertEqual(len(store), 1)
|
|
|
|
store.clear()
|
|
self.assertEqual(len(store), 0)
|
|
|
|
def test_export_json(self):
|
|
"""Export should return valid JSON."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
store.track({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
|
|
json_str = store.export_json()
|
|
data = json.loads(json_str)
|
|
self.assertEqual(len(data), 1)
|
|
self.assertEqual(data[0]["category"], "feature")
|
|
|
|
def test_rate_limiting(self):
|
|
"""Should reject events when rate limited."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
# Set very low rate limit
|
|
store._rate_limiter = RateLimiter(max_per_minute=2)
|
|
|
|
# First two should succeed
|
|
success1, _ = store.track({"category": "session", "action": "start"})
|
|
success2, _ = store.track({"category": "session", "action": "start"})
|
|
# Third should fail
|
|
success3, msg = store.track({"category": "session", "action": "start"})
|
|
|
|
self.assertTrue(success1)
|
|
self.assertTrue(success2)
|
|
self.assertFalse(success3)
|
|
self.assertIn("rate", msg.lower())
|
|
|
|
def test_pii_rejected(self):
|
|
"""Events with PII in label should be rejected."""
|
|
store = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
|
|
# Note: PII check happens in validate_event, but we need a valid category/action
|
|
# that allows arbitrary labels. Use analysis/pattern_matched with patched allowlist
|
|
with patch('telemetry.get_pattern_allowlist', return_value={"test@example.com"}):
|
|
# This should still fail due to PII detection
|
|
success, msg = store.track({
|
|
"category": "session",
|
|
"action": "start",
|
|
"label": "test@example.com", # Email in label
|
|
})
|
|
|
|
# Session labels are not validated by allowlist, but should catch PII
|
|
self.assertFalse(success)
|
|
self.assertIn("pii", msg.lower())
|
|
|
|
|
|
class TestTelemetryPersistence(unittest.TestCase):
|
|
"""Test telemetry file persistence."""
|
|
|
|
def setUp(self):
|
|
"""Create temp file for each test."""
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.temp_file = os.path.join(self.temp_dir, "telemetry.json")
|
|
|
|
def tearDown(self):
|
|
"""Clean up temp files."""
|
|
if os.path.exists(self.temp_file):
|
|
os.remove(self.temp_file)
|
|
if os.path.exists(self.temp_dir):
|
|
os.rmdir(self.temp_dir)
|
|
|
|
def test_persistence_across_instances(self):
|
|
"""Events should persist across store instances."""
|
|
# First instance
|
|
store1 = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
store1.track({
|
|
"category": "feature",
|
|
"action": "tab_switch",
|
|
"label": "chat",
|
|
})
|
|
self.assertEqual(len(store1), 1)
|
|
|
|
# Second instance (new object, same file)
|
|
store2 = TelemetryStore(filepath=self.temp_file, enabled=True)
|
|
self.assertEqual(len(store2), 1)
|
|
events = store2.get_buffer()
|
|
self.assertEqual(events[0]["label"], "chat")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|