53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from src.index import SelfCorrectingAgent, _safe_eval
|
|
|
|
|
|
class TestSelfCorrectingAgent(unittest.TestCase):
|
|
def setUp(self):
|
|
# Create a temporary file for knowledge persistence
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.knowledge_file = Path(self.temp_dir.name) / "knowledge.json"
|
|
self.agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_safe_eval_basic(self):
|
|
self.assertEqual(_safe_eval("2+3*4"), 14)
|
|
self.assertAlmostEqual(_safe_eval("10/4"), 2.5)
|
|
self.assertEqual(_safe_eval("-5 + 2"), -3)
|
|
|
|
def test_safe_eval_invalid(self):
|
|
with self.assertRaises(ValueError):
|
|
_safe_eval("import os; os.system('echo hi')")
|
|
with self.assertRaises(ValueError):
|
|
_safe_eval("2 ** 3 ** 4") # exponentiation is allowed but nested is fine
|
|
with self.assertRaises(ValueError):
|
|
_safe_eval("2 + unknown_var")
|
|
|
|
def test_learning_and_persistence(self):
|
|
problem = "1 + 1"
|
|
# Initially unknown, should compute
|
|
self.assertEqual(self.agent.solve(problem), 2)
|
|
# Simulate user correction
|
|
self.agent.knowledge[problem] = 3
|
|
# Now should return learned answer
|
|
self.assertEqual(self.agent.solve(problem), 3)
|
|
# Persist knowledge
|
|
self.agent._save_knowledge()
|
|
# Load into new agent
|
|
new_agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)
|
|
self.assertEqual(new_agent.solve(problem), 3)
|
|
|
|
def test_invalid_expression(self):
|
|
with self.assertRaises(ValueError):
|
|
self.agent.solve("2 + * 3")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |