ai-agent/tests/memory/test_audit.py

116 lines
No EOL
4.3 KiB
Python

import unittest
import threading
import time
from datetime import datetime, timedelta
from security.memory.audit import MemoryAudit
from security.rbac_engine import RBACEngine # Mock will be used in actual tests
class TestMemoryAudit(unittest.TestCase):
def setUp(self):
self.rbac = RBACEngine() # In real tests, use mock
self.audit = MemoryAudit(self.rbac)
def test_basic_logging(self):
"""Test basic audit logging functionality"""
self.audit.log("create", "user1", "resource1", {"key": "value"})
entries = self.audit.get_entries()
self.assertEqual(len(entries), 1)
self.assertEqual(entries[0]["operation"], "create")
self.assertEqual(entries[0]["user"], "user1")
def test_thread_safety(self):
"""Test that logging is thread-safe"""
def worker():
for i in range(100):
self.audit.log("update", f"user{i}", f"res{i}", {"data": i})
threads = [threading.Thread(target=worker) for _ in range(10)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(len(self.audit.get_entries()), 1000)
def test_filter_by_operation(self):
"""Test filtering by operation type"""
self.audit.log("create", "user1", "res1", {})
self.audit.log("update", "user2", "res2", {})
self.audit.log("delete", "user3", "res3", {})
creates = self.audit.by_operation("create")
self.assertEqual(len(creates), 1)
self.assertEqual(creates[0]["operation"], "create")
def test_filter_by_user(self):
"""Test filtering by user"""
self.audit.log("read", "alice", "doc1", {})
self.audit.log("read", "bob", "doc2", {})
alice_entries = self.audit.by_user("alice")
self.assertEqual(len(alice_entries), 1)
self.assertEqual(alice_entries[0]["user"], "alice")
def test_filter_by_time_range(self):
"""Test filtering by time range"""
now = datetime.utcnow()
past = now - timedelta(minutes=5)
future = now + timedelta(minutes=5)
# Log with controlled timestamps
with unittest.mock.patch('datetime.datetime') as mock_datetime:
mock_datetime.utcnow.return_value = past
self.audit.log("old", "user", "res", {})
mock_datetime.utcnow.return_value = now
self.audit.log("current", "user", "res", {})
mock_datetime.utcnow.return_value = future
self.audit.log("future", "user", "res", {})
# Get entries between past and now
recent = self.audit.by_time_range(
past.isoformat(),
now.isoformat()
)
self.assertEqual(len(recent), 2) # old and current
def test_integrity_check(self):
"""Test log integrity verification"""
self.audit.log("create", "user", "res", {})
self.audit.log("update", "user", "res", {})
hashes = self.audit.get_entry_hashes()
self.assertTrue(self.audit.verify_integrity(hashes))
def test_rbac_enforcement(self):
"""Test RBAC enforcement during audit operations"""
# Mock RBAC to allow access
mock_rbac = unittest.mock.MagicMock()
mock_rbac.check_access.return_value = (True, "Access granted")
audit = MemoryAudit(mock_rbac)
# Should succeed when RBAC allows
audit.log("read", "allowed_user", "resource", {})
entries = audit.get_entries()
self.assertEqual(len(entries), 1)
# Mock RBAC to deny access
mock_rbac.check_access.return_value = (False, "Access denied")
with self.assertRaises(PermissionError):
audit.log("read", "denied_user", "resource", {})
def test_audit_logs_access_decisions(self):
"""Test that access decisions are properly logged"""
mock_rbac = unittest.mock.MagicMock()
mock_rbac.check_access.return_value = (True, "Access granted")
audit = MemoryAudit(mock_rbac)
audit.log("read", "test_user", "test_resource", {})
entries = audit.get_entries()
self.assertIn("access_decision", entries[0])
self.assertEqual(entries[0]["access_decision"], "Access granted")
if __name__ == '__main__':
unittest.main()