Repository navigation
Expand file tree
/
Copy pathtest_telemetry.py
More file actions
144 lines (115 loc) · 4.74 KB
/
Copy pathtest_telemetry.py
File metadata and controls
144 lines (115 loc) · 4.74 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""Tests for Logfire telemetry bootstrap."""
from __future__ import annotations
from basic_memory import __version__, telemetry
from basic_memory.config import init_api_logging, init_cli_logging, init_mcp_logging
from typing import Any
class FakeLogfire:
"""Minimal fake of the logfire module surface configure_telemetry touches."""
class CodeSource:
def __init__(self, **kwargs):
self.kwargs = kwargs
def __init__(self, *, fail_on_send_to_logfire: bool = False) -> None:
self.fail_on_send_to_logfire = fail_on_send_to_logfire
self.configure_calls: list[dict[str, Any]] = []
def configure(self, **kwargs) -> None:
self.configure_calls.append(kwargs)
if self.fail_on_send_to_logfire and "send_to_logfire" in kwargs:
raise TypeError("send_to_logfire not supported")
def loguru_handler(self) -> dict[str, Any]:
return {"sink": "fake-logfire", "level": "INFO"}
def test_configure_telemetry_disabled_is_noop() -> None:
enabled = telemetry.configure_telemetry(
"basic-memory-cli",
environment="dev",
enable_logfire=False,
)
assert enabled is False
assert telemetry.get_logfire_handler() is None
def test_configure_telemetry_configures_logfire(monkeypatch) -> None:
fake_logfire = FakeLogfire()
monkeypatch.setattr(telemetry, "logfire", fake_logfire)
enabled = telemetry.configure_telemetry(
"basic-memory-api",
environment="prod",
service_version="abc123",
enable_logfire=True,
send_to_logfire=True,
)
assert enabled is True
assert telemetry.get_logfire_handler() == {"sink": "fake-logfire", "level": "INFO"}
assert len(fake_logfire.configure_calls) == 1
assert fake_logfire.configure_calls[0]["send_to_logfire"] is True
def test_configure_telemetry_retries_without_send_to_logfire(monkeypatch) -> None:
fake_logfire = FakeLogfire(fail_on_send_to_logfire=True)
monkeypatch.setattr(telemetry, "logfire", fake_logfire)
enabled = telemetry.configure_telemetry(
"basic-memory-api",
environment="prod",
enable_logfire=True,
send_to_logfire=True,
)
assert enabled is True
assert len(fake_logfire.configure_calls) == 2
assert fake_logfire.configure_calls[0]["send_to_logfire"] is True
assert "send_to_logfire" not in fake_logfire.configure_calls[1]
def test_configure_telemetry_clears_handler_when_disabled(monkeypatch) -> None:
"""A prior enabled configure must not leak its handler into a later disabled one."""
fake_logfire = FakeLogfire()
monkeypatch.setattr(telemetry, "logfire", fake_logfire)
telemetry.configure_telemetry("basic-memory-cli", environment="dev", enable_logfire=True)
assert telemetry.get_logfire_handler() is not None
telemetry.configure_telemetry("basic-memory-cli", environment="dev", enable_logfire=False)
assert telemetry.get_logfire_handler() is None
def test_init_logging_functions_configure_telemetry_and_logging(monkeypatch) -> None:
telemetry_calls: list[dict[str, Any]] = []
setup_calls: list[dict[str, Any]] = []
class StubConfig:
logfire_enabled = True
logfire_send_to_logfire = False
logfire_service_name = "basic-memory"
logfire_environment = "staging"
env = "dev"
monkeypatch.setattr(
"basic_memory.config.ConfigManager", lambda: type("CM", (), {"config": StubConfig()})()
)
monkeypatch.setattr(
"basic_memory.config.configure_telemetry",
lambda **kwargs: telemetry_calls.append(kwargs),
)
monkeypatch.setattr(
"basic_memory.config.setup_logging",
lambda **kwargs: setup_calls.append(kwargs),
)
monkeypatch.setenv("BASIC_MEMORY_LOG_LEVEL", "DEBUG")
monkeypatch.delenv("BASIC_MEMORY_CLOUD_MODE", raising=False)
init_cli_logging()
init_mcp_logging()
init_api_logging()
assert telemetry_calls == [
{
"service_name": "basic-memory-cli",
"environment": "staging",
"service_version": __version__,
"enable_logfire": True,
"send_to_logfire": False,
},
{
"service_name": "basic-memory-mcp",
"environment": "staging",
"service_version": __version__,
"enable_logfire": True,
"send_to_logfire": False,
},
{
"service_name": "basic-memory-api",
"environment": "staging",
"service_version": __version__,
"enable_logfire": True,
"send_to_logfire": False,
},
]
assert setup_calls == [
{"log_level": "DEBUG", "log_to_file": True},
{"log_level": "DEBUG", "log_to_file": True},
{"log_level": "DEBUG", "log_to_file": True},
]