25d54dd896
* auto fixes * auto fix pt2 and transitive deps and undefined var checking locals() * manual fixes (ignored or letta-code fixed) * fix circular import
152 lines
4.6 KiB
Python
152 lines
4.6 KiB
Python
import json
|
|
import logging
|
|
from io import StringIO
|
|
|
|
from letta.log import JSONFormatter, LogContextFilter
|
|
from letta.log_context import clear_log_context, get_log_context, remove_log_context, set_log_context, update_log_context
|
|
|
|
|
|
class TestLogContext:
|
|
def test_set_log_context(self):
|
|
clear_log_context()
|
|
set_log_context("agent_id", "agent-123")
|
|
assert get_log_context("agent_id") == "agent-123"
|
|
clear_log_context()
|
|
|
|
def test_update_log_context(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
context = get_log_context()
|
|
assert context["agent_id"] == "agent-123"
|
|
assert context["actor_id"] == "user-456"
|
|
clear_log_context()
|
|
|
|
def test_remove_log_context(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
remove_log_context("agent_id")
|
|
context = get_log_context()
|
|
assert "agent_id" not in context
|
|
assert context["actor_id"] == "user-456"
|
|
clear_log_context()
|
|
|
|
def test_clear_log_context(self):
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
clear_log_context()
|
|
context = get_log_context()
|
|
assert context == {}
|
|
|
|
def test_get_log_context_all(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
context = get_log_context()
|
|
assert isinstance(context, dict)
|
|
assert len(context) == 2
|
|
clear_log_context()
|
|
|
|
|
|
class TestLogContextFilter:
|
|
def test_filter_adds_context_to_record(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
|
|
log_filter = LogContextFilter()
|
|
record = logging.LogRecord(
|
|
name="test",
|
|
level=logging.INFO,
|
|
pathname="test.py",
|
|
lineno=1,
|
|
msg="Test message",
|
|
args=(),
|
|
exc_info=None,
|
|
)
|
|
|
|
result = log_filter.filter(record)
|
|
assert result is True
|
|
assert hasattr(record, "agent_id")
|
|
assert record.agent_id == "agent-123"
|
|
assert hasattr(record, "actor_id")
|
|
assert record.actor_id == "user-456"
|
|
clear_log_context()
|
|
|
|
def test_filter_does_not_override_existing_attributes(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123")
|
|
|
|
log_filter = LogContextFilter()
|
|
record = logging.LogRecord(
|
|
name="test",
|
|
level=logging.INFO,
|
|
pathname="test.py",
|
|
lineno=1,
|
|
msg="Test message",
|
|
args=(),
|
|
exc_info=None,
|
|
)
|
|
record.agent_id = "agent-999"
|
|
|
|
log_filter.filter(record)
|
|
assert record.agent_id == "agent-999"
|
|
clear_log_context()
|
|
|
|
|
|
class TestLogContextIntegration:
|
|
def test_json_formatter_includes_context(self):
|
|
clear_log_context()
|
|
update_log_context(agent_id="agent-123", actor_id="user-456")
|
|
|
|
logger = logging.getLogger("test_logger")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
handler = logging.StreamHandler(StringIO())
|
|
handler.setFormatter(JSONFormatter())
|
|
handler.addFilter(LogContextFilter())
|
|
logger.addHandler(handler)
|
|
|
|
log_stream = handler.stream
|
|
logger.info("Test message")
|
|
|
|
log_stream.seek(0)
|
|
log_output = log_stream.read()
|
|
|
|
log_data = json.loads(log_output)
|
|
assert log_data["message"] == "Test message"
|
|
assert log_data["agent_id"] == "agent-123"
|
|
assert log_data["actor_id"] == "user-456"
|
|
|
|
logger.removeHandler(handler)
|
|
clear_log_context()
|
|
|
|
def test_multiple_log_calls_with_changing_context(self):
|
|
clear_log_context()
|
|
logger = logging.getLogger("test_logger_2")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
handler = logging.StreamHandler(StringIO())
|
|
handler.setFormatter(JSONFormatter())
|
|
handler.addFilter(LogContextFilter())
|
|
logger.addHandler(handler)
|
|
|
|
log_stream = handler.stream
|
|
|
|
update_log_context(agent_id="agent-123")
|
|
logger.info("First message")
|
|
|
|
update_log_context(actor_id="user-456")
|
|
logger.info("Second message")
|
|
|
|
log_stream.seek(0)
|
|
log_lines = log_stream.readlines()
|
|
assert len(log_lines) == 2
|
|
|
|
first_log = json.loads(log_lines[0])
|
|
assert first_log["agent_id"] == "agent-123"
|
|
assert "actor_id" not in first_log
|
|
|
|
second_log = json.loads(log_lines[1])
|
|
assert second_log["agent_id"] == "agent-123"
|
|
assert second_log["actor_id"] == "user-456"
|
|
|
|
logger.removeHandler(handler)
|
|
clear_log_context()
|