96 lines
3 KiB
Python
96 lines
3 KiB
Python
"""
|
|
Test visualization functionality.
|
|
|
|
Tests memory graph data retrieval (not actual rendering).
|
|
"""
|
|
import pytest
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
|
|
def utcnow():
|
|
"""Get current UTC time with timezone info."""
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def test_get_memory_graph_data(memory, clean_agent):
|
|
"""Test retrieval of memory graph data for visualization."""
|
|
agent_id = clean_agent
|
|
|
|
# Store some memories
|
|
memory.put(
|
|
agent_id=agent_id,
|
|
content="Alice loves hiking in the mountains.",
|
|
context="Hobbies",
|
|
event_date=utcnow() - timedelta(hours=2),
|
|
)
|
|
|
|
memory.put(
|
|
agent_id=agent_id,
|
|
content="Bob enjoys rock climbing.",
|
|
context="Sports",
|
|
event_date=utcnow() - timedelta(hours=1),
|
|
)
|
|
|
|
memory.put(
|
|
agent_id=agent_id,
|
|
content="Alice is working on a Python project.",
|
|
context="Tech",
|
|
event_date=utcnow(),
|
|
)
|
|
|
|
# Get graph data
|
|
units, links = memory.get_memory_graph_data(agent_id)
|
|
|
|
assert isinstance(units, list), "Units should be a list"
|
|
assert isinstance(links, list), "Links should be a list"
|
|
assert len(units) > 0, "Should have memory units"
|
|
|
|
# Verify unit structure
|
|
for unit in units:
|
|
assert 'id' in unit, "Unit should have id"
|
|
assert 'text' in unit, "Unit should have text"
|
|
assert 'context' in unit, "Unit should have context"
|
|
assert 'event_date' in unit, "Unit should have event_date"
|
|
assert 'access_count' in unit, "Unit should have access_count"
|
|
|
|
# Links may or may not exist depending on similarity/proximity
|
|
if len(links) > 0:
|
|
# Verify link structure
|
|
for link in links:
|
|
assert 'from_unit_id' in link, "Link should have from_unit_id"
|
|
assert 'to_unit_id' in link, "Link should have to_unit_id"
|
|
assert 'link_type' in link, "Link should have link_type"
|
|
assert 'weight' in link, "Link should have weight"
|
|
assert link['link_type'] in ['temporal', 'semantic', 'entity'], \
|
|
"Link type should be temporal, semantic, or entity"
|
|
|
|
|
|
def test_memory_graph_has_correct_agent_data(memory, clean_agent):
|
|
"""Test that graph data only includes data for the specified agent."""
|
|
agent_id = clean_agent
|
|
other_agent_id = "other_agent"
|
|
|
|
# Store memories for test agent
|
|
memory.put(
|
|
agent_id=agent_id,
|
|
content="Alice loves hiking.",
|
|
context="Hobbies",
|
|
event_date=utcnow(),
|
|
)
|
|
|
|
# Store memories for another agent
|
|
memory.put(
|
|
agent_id=other_agent_id,
|
|
content="Charlie enjoys swimming.",
|
|
context="Sports",
|
|
event_date=utcnow(),
|
|
)
|
|
|
|
# Get graph data for test agent
|
|
units, links = memory.get_memory_graph_data(agent_id)
|
|
|
|
# Should only include test agent's data
|
|
for unit in units:
|
|
# Verify by checking text content (Alice should be present, Charlie should not)
|
|
unit_text = unit['text']
|
|
assert 'Charlie' not in unit_text, "Should not include other agent's memories"
|