326 lines
11 KiB
Python
326 lines
11 KiB
Python
import pytest
|
|
from ra_aid.tools.memory import (
|
|
_global_memory,
|
|
get_memory_value,
|
|
emit_research_subtask,
|
|
emit_key_facts,
|
|
delete_key_facts,
|
|
emit_key_snippets,
|
|
delete_key_snippets,
|
|
emit_related_files,
|
|
get_related_files
|
|
)
|
|
|
|
@pytest.fixture
|
|
def reset_memory():
|
|
"""Reset global memory before each test"""
|
|
_global_memory['key_facts'] = {}
|
|
_global_memory['key_fact_id_counter'] = 0
|
|
_global_memory['key_snippets'] = {}
|
|
_global_memory['key_snippet_id_counter'] = 0
|
|
_global_memory['research_notes'] = []
|
|
_global_memory['plans'] = []
|
|
_global_memory['tasks'] = []
|
|
_global_memory['research_subtasks'] = []
|
|
_global_memory['related_files'] = set()
|
|
yield
|
|
# Clean up after test
|
|
_global_memory['key_facts'] = {}
|
|
_global_memory['key_fact_id_counter'] = 0
|
|
_global_memory['key_snippets'] = {}
|
|
_global_memory['key_snippet_id_counter'] = 0
|
|
_global_memory['research_notes'] = []
|
|
_global_memory['plans'] = []
|
|
_global_memory['tasks'] = []
|
|
_global_memory['research_subtasks'] = []
|
|
|
|
def test_emit_key_facts_single_fact(reset_memory):
|
|
"""Test emitting a single key fact using emit_key_facts"""
|
|
# Test with single fact
|
|
result = emit_key_facts.invoke({"facts": ["First fact"]})
|
|
assert result == "Facts stored."
|
|
assert _global_memory['key_facts'][0] == "First fact"
|
|
assert _global_memory['key_fact_id_counter'] == 1
|
|
|
|
def test_delete_key_facts_single_fact(reset_memory):
|
|
"""Test deleting a single key fact using delete_key_facts"""
|
|
# Add a fact
|
|
emit_key_facts.invoke({"facts": ["Test fact"]})
|
|
|
|
# Delete the fact
|
|
result = delete_key_facts.invoke({"fact_ids": [0]})
|
|
assert result == "Facts deleted."
|
|
assert 0 not in _global_memory['key_facts']
|
|
|
|
def test_delete_key_facts_invalid(reset_memory):
|
|
"""Test deleting non-existent facts returns empty list"""
|
|
# Try to delete non-existent fact
|
|
result = delete_key_facts.invoke({"fact_ids": [999]})
|
|
assert result == "Facts deleted."
|
|
|
|
# Add and delete a fact, then try to delete it again
|
|
emit_key_facts.invoke({"facts": ["Test fact"]})
|
|
delete_key_facts.invoke({"fact_ids": [0]})
|
|
result = delete_key_facts.invoke({"fact_ids": [0]})
|
|
assert result == "Facts deleted."
|
|
|
|
def test_get_memory_value_key_facts(reset_memory):
|
|
"""Test get_memory_value with key facts dictionary"""
|
|
# Empty key facts should return empty string
|
|
assert get_memory_value('key_facts') == ""
|
|
|
|
# Add some facts
|
|
emit_key_facts.invoke({"facts": ["First fact", "Second fact"]})
|
|
|
|
# Should return markdown formatted list
|
|
expected = "## 🔑 Key Fact #0\n\nFirst fact\n\n## 🔑 Key Fact #1\n\nSecond fact"
|
|
assert get_memory_value('key_facts') == expected
|
|
|
|
def test_get_memory_value_other_types(reset_memory):
|
|
"""Test get_memory_value remains compatible with other memory types"""
|
|
# Add some research notes
|
|
_global_memory['research_notes'].append("Note 1")
|
|
_global_memory['research_notes'].append("Note 2")
|
|
|
|
assert get_memory_value('research_notes') == "Note 1\nNote 2"
|
|
|
|
# Test with empty list
|
|
assert get_memory_value('plans') == ""
|
|
|
|
# Test with non-existent key
|
|
assert get_memory_value('nonexistent') == ""
|
|
|
|
def test_emit_key_facts(reset_memory):
|
|
"""Test emitting multiple key facts at once"""
|
|
# Test emitting multiple facts
|
|
facts = ["First fact", "Second fact", "Third fact"]
|
|
result = emit_key_facts.invoke({"facts": facts})
|
|
|
|
# Verify return message
|
|
assert result == "Facts stored."
|
|
|
|
# Verify facts stored in memory with correct IDs
|
|
assert _global_memory['key_facts'][0] == "First fact"
|
|
assert _global_memory['key_facts'][1] == "Second fact"
|
|
assert _global_memory['key_facts'][2] == "Third fact"
|
|
|
|
# Verify counter incremented correctly
|
|
assert _global_memory['key_fact_id_counter'] == 3
|
|
|
|
def test_delete_key_facts(reset_memory):
|
|
"""Test deleting multiple key facts"""
|
|
# Add some test facts
|
|
emit_key_facts.invoke({"facts": ["First fact", "Second fact", "Third fact"]})
|
|
|
|
# Test deleting mix of existing and non-existing IDs
|
|
result = delete_key_facts.invoke({"fact_ids": [0, 1, 999]})
|
|
|
|
# Verify success message
|
|
assert result == "Facts deleted."
|
|
|
|
# Verify correct facts removed from memory
|
|
assert 0 not in _global_memory['key_facts']
|
|
assert 1 not in _global_memory['key_facts']
|
|
assert 2 in _global_memory['key_facts'] # ID 2 should remain
|
|
assert _global_memory['key_facts'][2] == "Third fact"
|
|
|
|
def test_emit_key_snippets(reset_memory):
|
|
"""Test emitting multiple code snippets at once"""
|
|
# Test snippets with and without descriptions
|
|
snippets = [
|
|
{
|
|
"filepath": "test.py",
|
|
"line_number": 10,
|
|
"snippet": "def test():\n pass",
|
|
"description": "Test function"
|
|
},
|
|
{
|
|
"filepath": "main.py",
|
|
"line_number": 20,
|
|
"snippet": "print('hello')",
|
|
"description": None
|
|
}
|
|
]
|
|
|
|
# Emit snippets
|
|
result = emit_key_snippets.invoke({"snippets": snippets})
|
|
|
|
# Verify return message
|
|
assert result == "Snippets stored."
|
|
|
|
# Verify snippets stored correctly
|
|
assert _global_memory['key_snippets'][0] == snippets[0]
|
|
assert _global_memory['key_snippets'][1] == snippets[1]
|
|
|
|
# Verify counter incremented correctly
|
|
assert _global_memory['key_snippet_id_counter'] == 2
|
|
|
|
def test_delete_key_snippets(reset_memory):
|
|
"""Test deleting multiple code snippets"""
|
|
# Add test snippets
|
|
snippets = [
|
|
{
|
|
"filepath": "test1.py",
|
|
"line_number": 1,
|
|
"snippet": "code1",
|
|
"description": None
|
|
},
|
|
{
|
|
"filepath": "test2.py",
|
|
"line_number": 2,
|
|
"snippet": "code2",
|
|
"description": None
|
|
},
|
|
{
|
|
"filepath": "test3.py",
|
|
"line_number": 3,
|
|
"snippet": "code3",
|
|
"description": None
|
|
}
|
|
]
|
|
emit_key_snippets.invoke({"snippets": snippets})
|
|
|
|
# Test deleting mix of valid and invalid IDs
|
|
result = delete_key_snippets.invoke({"snippet_ids": [0, 1, 999]})
|
|
|
|
# Verify success message
|
|
assert result == "Snippets deleted."
|
|
|
|
# Verify correct snippets removed
|
|
assert 0 not in _global_memory['key_snippets']
|
|
assert 1 not in _global_memory['key_snippets']
|
|
assert 2 in _global_memory['key_snippets']
|
|
assert _global_memory['key_snippets'][2]['filepath'] == "test3.py"
|
|
|
|
def test_delete_key_snippets_empty(reset_memory):
|
|
"""Test deleting snippets with empty ID list"""
|
|
# Add a test snippet
|
|
snippet = {
|
|
"filepath": "test.py",
|
|
"line_number": 1,
|
|
"snippet": "code",
|
|
"description": None
|
|
}
|
|
emit_key_snippets.invoke({"snippets": [snippet]})
|
|
|
|
# Test with empty list
|
|
result = delete_key_snippets.invoke({"snippet_ids": []})
|
|
assert result == "Snippets deleted."
|
|
|
|
# Verify snippet still exists
|
|
assert 0 in _global_memory['key_snippets']
|
|
|
|
def test_emit_related_files_basic(reset_memory):
|
|
"""Test basic adding of files"""
|
|
# Test adding single file
|
|
result = emit_related_files.invoke({"files": ["test.py"]})
|
|
assert result == "Files noted."
|
|
assert get_related_files() == {"test.py"}
|
|
|
|
# Test adding multiple files
|
|
result = emit_related_files.invoke({"files": ["main.py", "utils.py"]})
|
|
assert result == "Files noted."
|
|
assert get_related_files() == {"test.py", "main.py", "utils.py"}
|
|
|
|
def test_get_related_files_empty(reset_memory):
|
|
"""Test getting related files when none added"""
|
|
assert get_related_files() == set()
|
|
|
|
def test_emit_related_files_duplicates(reset_memory):
|
|
"""Test that duplicate files are handled correctly"""
|
|
# Add initial files
|
|
result = emit_related_files.invoke({"files": ["test.py", "main.py"]})
|
|
assert result == "Files noted."
|
|
assert get_related_files() == {"test.py", "main.py"}
|
|
|
|
# Try adding duplicates
|
|
result = emit_related_files.invoke({"files": ["test.py", "main.py", "test.py"]})
|
|
assert result == "Files noted."
|
|
# Set should still only contain unique entries
|
|
assert get_related_files() == {"test.py", "main.py"}
|
|
|
|
def test_key_snippets_integration(reset_memory):
|
|
"""Integration test for key snippets functionality"""
|
|
# Initial snippets to add
|
|
snippets = [
|
|
{
|
|
"filepath": "file1.py",
|
|
"line_number": 10,
|
|
"snippet": "def func1():\n pass",
|
|
"description": "First function"
|
|
},
|
|
{
|
|
"filepath": "file2.py",
|
|
"line_number": 20,
|
|
"snippet": "def func2():\n return True",
|
|
"description": "Second function"
|
|
},
|
|
{
|
|
"filepath": "file3.py",
|
|
"line_number": 30,
|
|
"snippet": "class TestClass:\n pass",
|
|
"description": "Test class"
|
|
}
|
|
]
|
|
|
|
# Add all snippets
|
|
result = emit_key_snippets.invoke({"snippets": snippets})
|
|
assert result == "Snippets stored."
|
|
assert _global_memory['key_snippet_id_counter'] == 3
|
|
# Verify related files were tracked
|
|
assert _global_memory['related_files'] == {"file1.py", "file2.py", "file3.py"}
|
|
|
|
# Verify all snippets were stored correctly
|
|
assert len(_global_memory['key_snippets']) == 3
|
|
assert _global_memory['key_snippets'][0] == snippets[0]
|
|
assert _global_memory['key_snippets'][1] == snippets[1]
|
|
assert _global_memory['key_snippets'][2] == snippets[2]
|
|
|
|
# Delete some but not all snippets (0 and 2)
|
|
result = delete_key_snippets.invoke({"snippet_ids": [0, 2]})
|
|
assert result == "Snippets deleted."
|
|
|
|
# Verify remaining snippet is intact
|
|
assert len(_global_memory['key_snippets']) == 1
|
|
assert 1 in _global_memory['key_snippets']
|
|
assert _global_memory['key_snippets'][1] == snippets[1]
|
|
|
|
# Counter should remain unchanged after deletions
|
|
assert _global_memory['key_snippet_id_counter'] == 3
|
|
|
|
# Add new snippet to verify counter continues correctly
|
|
new_snippet = {
|
|
"filepath": "file4.py",
|
|
"line_number": 40,
|
|
"snippet": "def func4():\n return False",
|
|
"description": "Fourth function"
|
|
}
|
|
result = emit_key_snippets.invoke({"snippets": [new_snippet]})
|
|
assert result == "Snippets stored."
|
|
assert _global_memory['key_snippet_id_counter'] == 4
|
|
# Verify new file was added to related files
|
|
assert _global_memory['related_files'] == {"file1.py", "file2.py", "file3.py", "file4.py"}
|
|
|
|
# Delete remaining snippets
|
|
result = delete_key_snippets.invoke({"snippet_ids": [1, 3]})
|
|
assert result == "Snippets deleted."
|
|
|
|
# Verify all snippets are gone
|
|
assert len(_global_memory['key_snippets']) == 0
|
|
|
|
# Counter should still maintain its value
|
|
assert _global_memory['key_snippet_id_counter'] == 4
|
|
|
|
def test_emit_research_subtask(reset_memory):
|
|
"""Test emitting research subtasks"""
|
|
# Test adding a research subtask
|
|
subtask = "Research Python async patterns"
|
|
result = emit_research_subtask(subtask)
|
|
|
|
# Verify return message
|
|
assert result == "Subtask added."
|
|
|
|
# Verify it was stored in memory
|
|
assert len(_global_memory['research_subtasks']) == 1
|
|
assert _global_memory['research_subtasks'][0] == subtask
|