cardosofelipe/syndarix
1
0
Fork 0
forked from cardosofelipe/fast-next-template
Code Issues 18 Pull Requests Wiki Activity

test(safety): add comprehensive tests for safety framework modules

Browse Source 
Add tests to improve backend coverage from 85% to 93%:

- test_audit.py: 60 tests for AuditLogger (20% -> 99%)
  - Hash chain integrity, sanitization, retention, handlers
  - Fixed bug: hash chain modification after event creation
  - Fixed bug: verification not using correct prev_hash

- test_hitl.py: Tests for HITL manager (0% -> 100%)
- test_permissions.py: Tests for permissions manager (0% -> 99%)
- test_rollback.py: Tests for rollback manager (0% -> 100%)
- test_metrics.py: Tests for metrics collector (0% -> 100%)
- test_mcp_integration.py: Tests for MCP safety wrapper (0% -> 100%)
- test_validation.py: Additional cache and edge case tests (76% -> 100%)
- test_scoring.py: Lock cleanup and edge case tests (78% -> 91%)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Felipe Cardoso
2026-01-04 19:41:54 +01:00
parent 758052dcff
commit 60ebeaa582
10 changed files with 6025 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
backend/app/services/context/adapters/claude.py
View File

@@ -123,7 +123,9 @@ class ClaudeAdapter(ModelAdapter):
if score:
# Escape score to prevent XML injection via metadata
escaped_score = self._escape_xml(str(score))
parts.append(f'<document source="{source}" relevance="{escaped_score}">')
parts.append(
f'<document source="{source}" relevance="{escaped_score}">'
)
else:
parts.append(f'<document source="{source}">')

36
backend/app/services/safety/audit/logger.py
View File

@@ -24,6 +24,9 @@ from ..models import (
logger = logging.getLogger(__name__)
# Sentinel for distinguishing "no argument passed" from "explicitly passing None"
_UNSET = object()
class AuditLogger:
"""
@@ -142,8 +145,10 @@ class AuditLogger:
# Add hash chain for tamper detection
if self._enable_hash_chain:
event_hash = self._compute_hash(event)
sanitized_details["_hash"] = event_hash
sanitized_details["_prev_hash"] = self._last_hash
# Modify event.details directly (not sanitized_details)
# to ensure the hash is stored on the actual event
event.details["_hash"] = event_hash
event.details["_prev_hash"] = self._last_hash
self._last_hash = event_hash
self._buffer.append(event)
@@ -415,7 +420,8 @@ class AuditLogger:
)
if stored_hash:
computed = self._compute_hash(event)
# Pass prev_hash to compute hash with correct chain position
computed = self._compute_hash(event, prev_hash=prev_hash)
if computed != stored_hash:
issues.append(
f"Hash mismatch at event {event.id}: "
@@ -462,9 +468,23 @@ class AuditLogger:
return sanitized
def _compute_hash(self, event: AuditEvent) -> str:
"""Compute hash for an event (excluding hash fields)."""
data = {
def _compute_hash(
self, event: AuditEvent, prev_hash: str | None | object = _UNSET
) -> str:
"""Compute hash for an event (excluding hash fields).
Args:
event: The audit event to hash.
prev_hash: Optional previous hash to use instead of self._last_hash.
Pass this during verification to use the correct chain.
Use None explicitly to indicate no previous hash.
"""
# Use passed prev_hash if explicitly provided, otherwise use instance state
effective_prev: str | None = (
self._last_hash if prev_hash is _UNSET else prev_hash # type: ignore[assignment]
)
data: dict[str, str | dict[str, str] | None] = {
"id": event.id,
"event_type": event.event_type.value,
"timestamp": event.timestamp.isoformat(),
@@ -480,8 +500,8 @@ class AuditLogger:
"correlation_id": event.correlation_id,
}
if self._last_hash:
data["_prev_hash"] = self._last_hash
if effective_prev:
data["_prev_hash"] = effective_prev
serialized = json.dumps(data, sort_keys=True, default=str)
return hashlib.sha256(serialized.encode()).hexdigest()

133
backend/tests/services/context/test_scoring.py
View File

@@ -758,3 +758,136 @@ class TestBaseScorer:
# Boundaries
assert scorer.normalize_score(0.0) == 0.0
assert scorer.normalize_score(1.0) == 1.0
class TestCompositeScorerEdgeCases:
"""Tests for CompositeScorer edge cases and lock management."""
@pytest.mark.asyncio
async def test_score_with_zero_weights(self) -> None:
"""Test scoring when all weights are zero."""
scorer = CompositeScorer(
relevance_weight=0.0,
recency_weight=0.0,
priority_weight=0.0,
)
context = KnowledgeContext(
content="Test content",
source="docs",
relevance_score=0.8,
)
# Should return 0.0 when total weight is 0
score = await scorer.score(context, "test query")
assert score == 0.0
@pytest.mark.asyncio
async def test_score_batch_sequential(self) -> None:
"""Test batch scoring in sequential mode (parallel=False)."""
scorer = CompositeScorer()
contexts = [
KnowledgeContext(
content="Content 1",
source="docs",
relevance_score=0.8,
),
KnowledgeContext(
content="Content 2",
source="docs",
relevance_score=0.5,
),
]
# Use parallel=False to cover the sequential path
scored = await scorer.score_batch(contexts, "query", parallel=False)
assert len(scored) == 2
assert scored[0].relevance_score == 0.8
assert scored[1].relevance_score == 0.5
@pytest.mark.asyncio
async def test_lock_fast_path_reuse(self) -> None:
"""Test that existing locks are reused via fast path."""
scorer = CompositeScorer()
context = KnowledgeContext(
content="Test",
source="docs",
relevance_score=0.5,
)
# First access creates the lock
lock1 = await scorer._get_context_lock(context.id)
# Second access should hit the fast path (lock exists in dict)
lock2 = await scorer._get_context_lock(context.id)
assert lock2 is lock1 # Same lock object returned
@pytest.mark.asyncio
async def test_lock_cleanup_when_limit_reached(self) -> None:
"""Test that old locks are cleaned up when limit is reached."""
import time
# Create scorer with very low max_locks to trigger cleanup
scorer = CompositeScorer()
scorer._max_locks = 3
scorer._lock_ttl = 0.1 # 100ms TTL
# Create locks for several context IDs
context_ids = [f"ctx-{i}" for i in range(5)]
# Get locks for first 3 contexts (fill up to limit)
for ctx_id in context_ids[:3]:
await scorer._get_context_lock(ctx_id)
# Wait for TTL to expire
time.sleep(0.15)
# Getting a lock for a new context should trigger cleanup
await scorer._get_context_lock(context_ids[3])
# Some old locks should have been cleaned up
# The exact number depends on cleanup logic
assert len(scorer._context_locks) <= scorer._max_locks + 1
@pytest.mark.asyncio
async def test_lock_cleanup_preserves_held_locks(self) -> None:
"""Test that cleanup doesn't remove locks that are currently held."""
import time
scorer = CompositeScorer()
scorer._max_locks = 2
scorer._lock_ttl = 0.05 # 50ms TTL
# Get and hold lock1
lock1 = await scorer._get_context_lock("ctx-1")
async with lock1:
# While holding lock1, add more locks
await scorer._get_context_lock("ctx-2")
time.sleep(0.1) # Let TTL expire
# Adding another should trigger cleanup
await scorer._get_context_lock("ctx-3")
# lock1 should still exist (it's held)
assert any(lock is lock1 for lock, _ in scorer._context_locks.values())
@pytest.mark.asyncio
async def test_concurrent_lock_acquisition_double_check(self) -> None:
"""Test that concurrent lock acquisition uses double-check pattern."""
import asyncio
scorer = CompositeScorer()
context_id = "test-context-id"
# Simulate concurrent lock acquisition
async def get_lock():
return await scorer._get_context_lock(context_id)
locks = await asyncio.gather(*[get_lock() for _ in range(10)])
# All should get the same lock (double-check pattern ensures this)
assert all(lock is locks[0] for lock in locks)

989
backend/tests/services/safety/test_audit.py Normal file
View File

@@ -0,0 +1,989 @@
"""
Tests for Audit Logger.
Tests cover:
- AuditLogger initialization and lifecycle
- Event logging and hash chain
- Query and filtering
- Retention policy enforcement
- Handler management
- Singleton pattern
"""
import asyncio
from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
import pytest
import pytest_asyncio
from app.services.safety.audit.logger import (
AuditLogger,
get_audit_logger,
reset_audit_logger,
shutdown_audit_logger,
)
from app.services.safety.models import (
ActionMetadata,
ActionRequest,
ActionType,
AuditEventType,
AutonomyLevel,
SafetyDecision,
)
class TestAuditLoggerInit:
"""Tests for AuditLogger initialization."""
def test_init_default_values(self):
"""Test initialization with default values."""
logger = AuditLogger()
assert logger._flush_interval == 10.0
assert logger._enable_hash_chain is True
assert logger._last_hash is None
assert logger._running is False
def test_init_custom_values(self):
"""Test initialization with custom values."""
logger = AuditLogger(
max_buffer_size=500,
flush_interval_seconds=5.0,
enable_hash_chain=False,
)
assert logger._flush_interval == 5.0
assert logger._enable_hash_chain is False
class TestAuditLoggerLifecycle:
"""Tests for AuditLogger start/stop."""
@pytest.mark.asyncio
async def test_start_creates_flush_task(self):
"""Test that start creates the periodic flush task."""
logger = AuditLogger(flush_interval_seconds=1.0)
await logger.start()
assert logger._running is True
assert logger._flush_task is not None
await logger.stop()
@pytest.mark.asyncio
async def test_start_idempotent(self):
"""Test that multiple starts don't create multiple tasks."""
logger = AuditLogger()
await logger.start()
task1 = logger._flush_task
await logger.start() # Second start
task2 = logger._flush_task
assert task1 is task2
await logger.stop()
@pytest.mark.asyncio
async def test_stop_cancels_task_and_flushes(self):
"""Test that stop cancels the task and flushes events."""
logger = AuditLogger()
await logger.start()
# Add an event
await logger.log(AuditEventType.ACTION_REQUESTED, agent_id="agent-1")
await logger.stop()
assert logger._running is False
# Event should be flushed
assert len(logger._persisted) == 1
@pytest.mark.asyncio
async def test_stop_without_start(self):
"""Test stopping without starting doesn't error."""
logger = AuditLogger()
await logger.stop() # Should not raise
class TestAuditLoggerLog:
"""Tests for the log method."""
@pytest_asyncio.fixture
async def logger(self):
"""Create a logger instance."""
return AuditLogger(enable_hash_chain=True)
@pytest.mark.asyncio
async def test_log_creates_event(self, logger):
"""Test logging creates an event."""
event = await logger.log(
AuditEventType.ACTION_REQUESTED,
agent_id="agent-1",
project_id="proj-1",
)
assert event.event_type == AuditEventType.ACTION_REQUESTED
assert event.agent_id == "agent-1"
assert event.project_id == "proj-1"
assert event.id is not None
assert event.timestamp is not None
@pytest.mark.asyncio
async def test_log_adds_hash_chain(self, logger):
"""Test logging adds hash chain."""
event = await logger.log(
AuditEventType.ACTION_REQUESTED,
agent_id="agent-1",
)
assert "_hash" in event.details
assert "_prev_hash" in event.details
assert event.details["_prev_hash"] is None # First event
@pytest.mark.asyncio
async def test_log_chain_links_events(self, logger):
"""Test hash chain links events."""
event1 = await logger.log(AuditEventType.ACTION_REQUESTED)
event2 = await logger.log(AuditEventType.ACTION_EXECUTED)
assert event2.details["_prev_hash"] == event1.details["_hash"]
@pytest.mark.asyncio
async def test_log_without_hash_chain(self):
"""Test logging without hash chain."""
logger = AuditLogger(enable_hash_chain=False)
event = await logger.log(AuditEventType.ACTION_REQUESTED)
assert "_hash" not in event.details
assert "_prev_hash" not in event.details
@pytest.mark.asyncio
async def test_log_with_all_fields(self, logger):
"""Test logging with all optional fields."""
event = await logger.log(
AuditEventType.ACTION_EXECUTED,
agent_id="agent-1",
action_id="action-1",
project_id="proj-1",
session_id="sess-1",
user_id="user-1",
decision=SafetyDecision.ALLOW,
details={"custom": "data"},
correlation_id="corr-1",
)
assert event.agent_id == "agent-1"
assert event.action_id == "action-1"
assert event.project_id == "proj-1"
assert event.session_id == "sess-1"
assert event.user_id == "user-1"
assert event.decision == SafetyDecision.ALLOW
assert event.details["custom"] == "data"
assert event.correlation_id == "corr-1"
@pytest.mark.asyncio
async def test_log_buffers_event(self, logger):
"""Test logging adds event to buffer."""
await logger.log(AuditEventType.ACTION_REQUESTED)
assert len(logger._buffer) == 1
class TestAuditLoggerConvenienceMethods:
"""Tests for convenience logging methods."""
@pytest_asyncio.fixture
async def logger(self):
"""Create a logger instance."""
return AuditLogger(enable_hash_chain=False)
@pytest_asyncio.fixture
def action(self):
"""Create a test action request."""
metadata = ActionMetadata(
agent_id="agent-1",
session_id="sess-1",
project_id="proj-1",
autonomy_level=AutonomyLevel.MILESTONE,
user_id="user-1",
correlation_id="corr-1",
)
return ActionRequest(
action_type=ActionType.FILE_WRITE,
tool_name="file_write",
arguments={"path": "/test.txt"},
resource="/test.txt",
metadata=metadata,
)
@pytest.mark.asyncio
async def test_log_action_request_allowed(self, logger, action):
"""Test logging allowed action request."""
event = await logger.log_action_request(
action,
SafetyDecision.ALLOW,
reasons=["Within budget"],
)
assert event.event_type == AuditEventType.ACTION_VALIDATED
assert event.decision == SafetyDecision.ALLOW
assert event.details["reasons"] == ["Within budget"]
@pytest.mark.asyncio
async def test_log_action_request_denied(self, logger, action):
"""Test logging denied action request."""
event = await logger.log_action_request(
action,
SafetyDecision.DENY,
reasons=["Rate limit exceeded"],
)
assert event.event_type == AuditEventType.ACTION_DENIED
assert event.decision == SafetyDecision.DENY
@pytest.mark.asyncio
async def test_log_action_executed_success(self, logger, action):
"""Test logging successful action execution."""
event = await logger.log_action_executed(
action,
success=True,
execution_time_ms=50.0,
)
assert event.event_type == AuditEventType.ACTION_EXECUTED
assert event.details["success"] is True
assert event.details["execution_time_ms"] == 50.0
assert event.details["error"] is None
@pytest.mark.asyncio
async def test_log_action_executed_failure(self, logger, action):
"""Test logging failed action execution."""
event = await logger.log_action_executed(
action,
success=False,
execution_time_ms=100.0,
error="File not found",
)
assert event.event_type == AuditEventType.ACTION_FAILED
assert event.details["success"] is False
assert event.details["error"] == "File not found"
@pytest.mark.asyncio
async def test_log_approval_event(self, logger, action):
"""Test logging approval event."""
event = await logger.log_approval_event(
AuditEventType.APPROVAL_GRANTED,
approval_id="approval-1",
action=action,
decided_by="admin",
reason="Approved by admin",
)
assert event.event_type == AuditEventType.APPROVAL_GRANTED
assert event.details["approval_id"] == "approval-1"
assert event.details["decided_by"] == "admin"
@pytest.mark.asyncio
async def test_log_budget_event(self, logger):
"""Test logging budget event."""
event = await logger.log_budget_event(
AuditEventType.BUDGET_WARNING,
agent_id="agent-1",
scope="daily",
current_usage=8000.0,
limit=10000.0,
unit="tokens",
)
assert event.event_type == AuditEventType.BUDGET_WARNING
assert event.details["scope"] == "daily"
assert event.details["usage_percent"] == 80.0
@pytest.mark.asyncio
async def test_log_budget_event_zero_limit(self, logger):
"""Test logging budget event with zero limit."""
event = await logger.log_budget_event(
AuditEventType.BUDGET_WARNING,
agent_id="agent-1",
scope="daily",
current_usage=100.0,
limit=0.0, # Zero limit
)
assert event.details["usage_percent"] == 0
@pytest.mark.asyncio
async def test_log_emergency_stop(self, logger):
"""Test logging emergency stop."""
event = await logger.log_emergency_stop(
stop_type="global",
triggered_by="admin",
reason="Security incident",
affected_agents=["agent-1", "agent-2"],
)
assert event.event_type == AuditEventType.EMERGENCY_STOP
assert event.details["stop_type"] == "global"
assert event.details["affected_agents"] == ["agent-1", "agent-2"]
class TestAuditLoggerFlush:
"""Tests for flush functionality."""
@pytest.mark.asyncio
async def test_flush_persists_events(self):
"""Test flush moves events to persisted storage."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(AuditEventType.ACTION_REQUESTED)
await logger.log(AuditEventType.ACTION_EXECUTED)
assert len(logger._buffer) == 2
assert len(logger._persisted) == 0
count = await logger.flush()
assert count == 2
assert len(logger._buffer) == 0
assert len(logger._persisted) == 2
@pytest.mark.asyncio
async def test_flush_empty_buffer(self):
"""Test flush with empty buffer."""
logger = AuditLogger()
count = await logger.flush()
assert count == 0
class TestAuditLoggerQuery:
"""Tests for query functionality."""
@pytest_asyncio.fixture
async def logger_with_events(self):
"""Create a logger with some test events."""
logger = AuditLogger(enable_hash_chain=False)
# Add various events
await logger.log(
AuditEventType.ACTION_REQUESTED,
agent_id="agent-1",
project_id="proj-1",
)
await logger.log(
AuditEventType.ACTION_EXECUTED,
agent_id="agent-1",
project_id="proj-1",
)
await logger.log(
AuditEventType.ACTION_DENIED,
agent_id="agent-2",
project_id="proj-2",
)
await logger.log(
AuditEventType.BUDGET_WARNING,
agent_id="agent-1",
project_id="proj-1",
user_id="user-1",
)
return logger
@pytest.mark.asyncio
async def test_query_all(self, logger_with_events):
"""Test querying all events."""
events = await logger_with_events.query()
assert len(events) == 4
@pytest.mark.asyncio
async def test_query_by_event_type(self, logger_with_events):
"""Test filtering by event type."""
events = await logger_with_events.query(
event_types=[AuditEventType.ACTION_REQUESTED]
)
assert len(events) == 1
assert events[0].event_type == AuditEventType.ACTION_REQUESTED
@pytest.mark.asyncio
async def test_query_by_agent_id(self, logger_with_events):
"""Test filtering by agent ID."""
events = await logger_with_events.query(agent_id="agent-1")
assert len(events) == 3
assert all(e.agent_id == "agent-1" for e in events)
@pytest.mark.asyncio
async def test_query_by_project_id(self, logger_with_events):
"""Test filtering by project ID."""
events = await logger_with_events.query(project_id="proj-2")
assert len(events) == 1
@pytest.mark.asyncio
async def test_query_by_user_id(self, logger_with_events):
"""Test filtering by user ID."""
events = await logger_with_events.query(user_id="user-1")
assert len(events) == 1
assert events[0].event_type == AuditEventType.BUDGET_WARNING
@pytest.mark.asyncio
async def test_query_with_limit(self, logger_with_events):
"""Test query with limit."""
events = await logger_with_events.query(limit=2)
assert len(events) == 2
@pytest.mark.asyncio
async def test_query_with_offset(self, logger_with_events):
"""Test query with offset."""
all_events = await logger_with_events.query()
offset_events = await logger_with_events.query(offset=2)
assert len(offset_events) == 2
assert offset_events[0] == all_events[2]
@pytest.mark.asyncio
async def test_query_by_time_range(self):
"""Test filtering by time range."""
logger = AuditLogger(enable_hash_chain=False)
now = datetime.utcnow()
await logger.log(AuditEventType.ACTION_REQUESTED)
# Query with start time
events = await logger.query(
start_time=now - timedelta(seconds=1),
end_time=now + timedelta(seconds=1),
)
assert len(events) == 1
@pytest.mark.asyncio
async def test_query_by_correlation_id(self):
"""Test filtering by correlation ID."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(
AuditEventType.ACTION_REQUESTED,
correlation_id="corr-123",
)
await logger.log(
AuditEventType.ACTION_EXECUTED,
correlation_id="corr-456",
)
events = await logger.query(correlation_id="corr-123")
assert len(events) == 1
assert events[0].correlation_id == "corr-123"
@pytest.mark.asyncio
async def test_query_combined_filters(self, logger_with_events):
"""Test combined filters."""
events = await logger_with_events.query(
agent_id="agent-1",
project_id="proj-1",
event_types=[
AuditEventType.ACTION_REQUESTED,
AuditEventType.ACTION_EXECUTED,
],
)
assert len(events) == 2
@pytest.mark.asyncio
async def test_get_action_history(self, logger_with_events):
"""Test get_action_history method."""
events = await logger_with_events.get_action_history("agent-1")
# Should only return action-related events
assert len(events) == 2
assert all(
e.event_type
in {AuditEventType.ACTION_REQUESTED, AuditEventType.ACTION_EXECUTED}
for e in events
)
class TestAuditLoggerIntegrity:
"""Tests for hash chain integrity verification."""
@pytest.mark.asyncio
async def test_verify_integrity_valid(self):
"""Test integrity verification with valid chain."""
logger = AuditLogger(enable_hash_chain=True)
await logger.log(AuditEventType.ACTION_REQUESTED)
await logger.log(AuditEventType.ACTION_EXECUTED)
is_valid, issues = await logger.verify_integrity()
assert is_valid is True
assert len(issues) == 0
@pytest.mark.asyncio
async def test_verify_integrity_disabled(self):
"""Test integrity verification when hash chain disabled."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(AuditEventType.ACTION_REQUESTED)
is_valid, issues = await logger.verify_integrity()
assert is_valid is True
assert len(issues) == 0
@pytest.mark.asyncio
async def test_verify_integrity_broken_chain(self):
"""Test integrity verification with broken chain."""
logger = AuditLogger(enable_hash_chain=True)
event1 = await logger.log(AuditEventType.ACTION_REQUESTED)
await logger.log(AuditEventType.ACTION_EXECUTED)
# Tamper with first event's hash
event1.details["_hash"] = "tampered_hash"
is_valid, issues = await logger.verify_integrity()
assert is_valid is False
assert len(issues) > 0
class TestAuditLoggerHandlers:
"""Tests for event handler management."""
@pytest.mark.asyncio
async def test_add_sync_handler(self):
"""Test adding synchronous handler."""
logger = AuditLogger(enable_hash_chain=False)
events_received = []
def handler(event):
events_received.append(event)
logger.add_handler(handler)
await logger.log(AuditEventType.ACTION_REQUESTED)
assert len(events_received) == 1
@pytest.mark.asyncio
async def test_add_async_handler(self):
"""Test adding async handler."""
logger = AuditLogger(enable_hash_chain=False)
events_received = []
async def handler(event):
events_received.append(event)
logger.add_handler(handler)
await logger.log(AuditEventType.ACTION_REQUESTED)
assert len(events_received) == 1
@pytest.mark.asyncio
async def test_remove_handler(self):
"""Test removing handler."""
logger = AuditLogger(enable_hash_chain=False)
events_received = []
def handler(event):
events_received.append(event)
logger.add_handler(handler)
await logger.log(AuditEventType.ACTION_REQUESTED)
logger.remove_handler(handler)
await logger.log(AuditEventType.ACTION_EXECUTED)
assert len(events_received) == 1
@pytest.mark.asyncio
async def test_handler_error_caught(self):
"""Test that handler errors are caught."""
logger = AuditLogger(enable_hash_chain=False)
def failing_handler(event):
raise ValueError("Handler error")
logger.add_handler(failing_handler)
# Should not raise
event = await logger.log(AuditEventType.ACTION_REQUESTED)
assert event is not None
class TestAuditLoggerSanitization:
"""Tests for sensitive data sanitization."""
@pytest.mark.asyncio
async def test_sanitize_sensitive_keys(self):
"""Test sanitization of sensitive keys."""
with patch("app.services.safety.audit.logger.get_safety_config") as mock_config:
mock_cfg = MagicMock()
mock_cfg.audit_retention_days = 30
mock_cfg.audit_include_sensitive = False
mock_config.return_value = mock_cfg
logger = AuditLogger(enable_hash_chain=False)
event = await logger.log(
AuditEventType.ACTION_EXECUTED,
details={
"password": "secret123",
"api_key": "key123",
"token": "token123",
"normal_field": "visible",
},
)
assert event.details["password"] == "[REDACTED]"
assert event.details["api_key"] == "[REDACTED]"
assert event.details["token"] == "[REDACTED]"
assert event.details["normal_field"] == "visible"
@pytest.mark.asyncio
async def test_sanitize_nested_dict(self):
"""Test sanitization of nested dictionaries."""
with patch("app.services.safety.audit.logger.get_safety_config") as mock_config:
mock_cfg = MagicMock()
mock_cfg.audit_retention_days = 30
mock_cfg.audit_include_sensitive = False
mock_config.return_value = mock_cfg
logger = AuditLogger(enable_hash_chain=False)
event = await logger.log(
AuditEventType.ACTION_EXECUTED,
details={
"config": {
"api_secret": "secret",
"name": "test",
}
},
)
assert event.details["config"]["api_secret"] == "[REDACTED]"
assert event.details["config"]["name"] == "test"
@pytest.mark.asyncio
async def test_include_sensitive_when_enabled(self):
"""Test sensitive data included when enabled."""
with patch("app.services.safety.audit.logger.get_safety_config") as mock_config:
mock_cfg = MagicMock()
mock_cfg.audit_retention_days = 30
mock_cfg.audit_include_sensitive = True
mock_config.return_value = mock_cfg
logger = AuditLogger(enable_hash_chain=False)
event = await logger.log(
AuditEventType.ACTION_EXECUTED,
details={"password": "secret123"},
)
assert event.details["password"] == "secret123"
class TestAuditLoggerRetention:
"""Tests for retention policy enforcement."""
@pytest.mark.asyncio
async def test_retention_removes_old_events(self):
"""Test that retention removes old events."""
with patch("app.services.safety.audit.logger.get_safety_config") as mock_config:
mock_cfg = MagicMock()
mock_cfg.audit_retention_days = 7
mock_cfg.audit_include_sensitive = False
mock_config.return_value = mock_cfg
logger = AuditLogger(enable_hash_chain=False)
# Add an old event directly to persisted
from app.services.safety.models import AuditEvent
old_event = AuditEvent(
id="old-event",
event_type=AuditEventType.ACTION_REQUESTED,
timestamp=datetime.utcnow() - timedelta(days=10),
details={},
)
logger._persisted.append(old_event)
# Add a recent event
await logger.log(AuditEventType.ACTION_EXECUTED)
# Flush will trigger retention enforcement
await logger.flush()
# Old event should be removed
assert len(logger._persisted) == 1
assert logger._persisted[0].id != "old-event"
@pytest.mark.asyncio
async def test_retention_keeps_recent_events(self):
"""Test that retention keeps recent events."""
with patch("app.services.safety.audit.logger.get_safety_config") as mock_config:
mock_cfg = MagicMock()
mock_cfg.audit_retention_days = 7
mock_cfg.audit_include_sensitive = False
mock_config.return_value = mock_cfg
logger = AuditLogger(enable_hash_chain=False)
await logger.log(AuditEventType.ACTION_REQUESTED)
await logger.log(AuditEventType.ACTION_EXECUTED)
await logger.flush()
assert len(logger._persisted) == 2
class TestAuditLoggerSingleton:
"""Tests for singleton pattern."""
@pytest.mark.asyncio
async def test_get_audit_logger_creates_instance(self):
"""Test get_audit_logger creates singleton."""
reset_audit_logger()
logger1 = await get_audit_logger()
logger2 = await get_audit_logger()
assert logger1 is logger2
await shutdown_audit_logger()
@pytest.mark.asyncio
async def test_shutdown_audit_logger(self):
"""Test shutdown_audit_logger stops and clears singleton."""
import app.services.safety.audit.logger as audit_module
reset_audit_logger()
_logger = await get_audit_logger()
await shutdown_audit_logger()
assert audit_module._audit_logger is None
def test_reset_audit_logger(self):
"""Test reset_audit_logger clears singleton."""
import app.services.safety.audit.logger as audit_module
audit_module._audit_logger = AuditLogger()
reset_audit_logger()
assert audit_module._audit_logger is None
class TestAuditLoggerPeriodicFlush:
"""Tests for periodic flush background task."""
@pytest.mark.asyncio
async def test_periodic_flush_runs(self):
"""Test periodic flush runs and flushes events."""
logger = AuditLogger(flush_interval_seconds=0.1, enable_hash_chain=False)
await logger.start()
# Log an event
await logger.log(AuditEventType.ACTION_REQUESTED)
assert len(logger._buffer) == 1
# Wait for periodic flush
await asyncio.sleep(0.15)
# Event should be flushed
assert len(logger._buffer) == 0
assert len(logger._persisted) == 1
await logger.stop()
@pytest.mark.asyncio
async def test_periodic_flush_handles_errors(self):
"""Test periodic flush handles errors gracefully."""
logger = AuditLogger(flush_interval_seconds=0.1)
await logger.start()
# Mock flush to raise an error
original_flush = logger.flush
async def failing_flush():
raise Exception("Flush error")
logger.flush = failing_flush
# Wait for flush attempt
await asyncio.sleep(0.15)
# Should still be running
assert logger._running is True
logger.flush = original_flush
await logger.stop()
class TestAuditLoggerLogging:
"""Tests for standard logger output."""
@pytest.mark.asyncio
async def test_log_warning_for_denied(self):
"""Test warning level for denied events."""
with patch("app.services.safety.audit.logger.logger") as mock_logger:
audit_logger = AuditLogger(enable_hash_chain=False)
await audit_logger.log(
AuditEventType.ACTION_DENIED,
agent_id="agent-1",
)
mock_logger.warning.assert_called()
@pytest.mark.asyncio
async def test_log_error_for_failed(self):
"""Test error level for failed events."""
with patch("app.services.safety.audit.logger.logger") as mock_logger:
audit_logger = AuditLogger(enable_hash_chain=False)
await audit_logger.log(
AuditEventType.ACTION_FAILED,
agent_id="agent-1",
)
mock_logger.error.assert_called()
@pytest.mark.asyncio
async def test_log_info_for_normal(self):
"""Test info level for normal events."""
with patch("app.services.safety.audit.logger.logger") as mock_logger:
audit_logger = AuditLogger(enable_hash_chain=False)
await audit_logger.log(
AuditEventType.ACTION_EXECUTED,
agent_id="agent-1",
)
mock_logger.info.assert_called()
class TestAuditLoggerEdgeCases:
"""Tests for edge cases."""
@pytest.mark.asyncio
async def test_log_with_none_details(self):
"""Test logging with None details."""
logger = AuditLogger(enable_hash_chain=False)
event = await logger.log(
AuditEventType.ACTION_REQUESTED,
details=None,
)
assert event.details == {}
@pytest.mark.asyncio
async def test_query_with_action_id(self):
"""Test querying by action ID."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(
AuditEventType.ACTION_REQUESTED,
action_id="action-1",
)
await logger.log(
AuditEventType.ACTION_EXECUTED,
action_id="action-2",
)
events = await logger.query(action_id="action-1")
assert len(events) == 1
assert events[0].action_id == "action-1"
@pytest.mark.asyncio
async def test_query_with_session_id(self):
"""Test querying by session ID."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(
AuditEventType.ACTION_REQUESTED,
session_id="sess-1",
)
await logger.log(
AuditEventType.ACTION_EXECUTED,
session_id="sess-2",
)
events = await logger.query(session_id="sess-1")
assert len(events) == 1
@pytest.mark.asyncio
async def test_query_includes_buffer_and_persisted(self):
"""Test query includes both buffer and persisted events."""
logger = AuditLogger(enable_hash_chain=False)
# Add event to buffer
await logger.log(AuditEventType.ACTION_REQUESTED)
# Flush to persisted
await logger.flush()
# Add another to buffer
await logger.log(AuditEventType.ACTION_EXECUTED)
# Query should return both
events = await logger.query()
assert len(events) == 2
@pytest.mark.asyncio
async def test_remove_nonexistent_handler(self):
"""Test removing handler that doesn't exist."""
logger = AuditLogger()
def handler(event):
pass
# Should not raise
logger.remove_handler(handler)
@pytest.mark.asyncio
async def test_query_time_filter_excludes_events(self):
"""Test time filters exclude events correctly."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(AuditEventType.ACTION_REQUESTED)
# Query with future start time
future = datetime.utcnow() + timedelta(hours=1)
events = await logger.query(start_time=future)
assert len(events) == 0
@pytest.mark.asyncio
async def test_query_end_time_filter(self):
"""Test end time filter."""
logger = AuditLogger(enable_hash_chain=False)
await logger.log(AuditEventType.ACTION_REQUESTED)
# Query with past end time
past = datetime.utcnow() - timedelta(hours=1)
events = await logger.query(end_time=past)
assert len(events) == 0

1136
backend/tests/services/safety/test_hitl.py Normal file
View File

File diff suppressed because it is too large Load Diff

874
backend/tests/services/safety/test_mcp_integration.py Normal file
View File

@@ -0,0 +1,874 @@
"""
Tests for MCP Safety Integration.
Tests cover:
- MCPToolCall and MCPToolResult data structures
- MCPSafetyWrapper: tool registration, execution, safety checks
- Tool classification and action type mapping
- SafeToolExecutor context manager
- Factory function create_mcp_wrapper
"""
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import pytest_asyncio
from app.services.safety.exceptions import EmergencyStopError
from app.services.safety.mcp.integration import (
MCPSafetyWrapper,
MCPToolCall,
MCPToolResult,
SafeToolExecutor,
create_mcp_wrapper,
)
from app.services.safety.models import (
ActionType,
AutonomyLevel,
SafetyDecision,
)
class TestMCPToolCall:
"""Tests for MCPToolCall dataclass."""
def test_tool_call_creation(self):
"""Test creating a tool call."""
call = MCPToolCall(
tool_name="file_read",
arguments={"path": "/tmp/test.txt"}, # noqa: S108
server_name="file-server",
project_id="proj-1",
context={"session_id": "sess-1"},
)
assert call.tool_name == "file_read"
assert call.arguments == {"path": "/tmp/test.txt"} # noqa: S108
assert call.server_name == "file-server"
assert call.project_id == "proj-1"
assert call.context == {"session_id": "sess-1"}
def test_tool_call_defaults(self):
"""Test tool call default values."""
call = MCPToolCall(
tool_name="test",
arguments={},
)
assert call.server_name is None
assert call.project_id is None
assert call.context == {}
class TestMCPToolResult:
"""Tests for MCPToolResult dataclass."""
def test_tool_result_success(self):
"""Test creating a successful result."""
result = MCPToolResult(
success=True,
result={"data": "test"},
safety_decision=SafetyDecision.ALLOW,
execution_time_ms=50.0,
)
assert result.success is True
assert result.result == {"data": "test"}
assert result.error is None
assert result.safety_decision == SafetyDecision.ALLOW
assert result.execution_time_ms == 50.0
def test_tool_result_failure(self):
"""Test creating a failed result."""
result = MCPToolResult(
success=False,
error="Permission denied",
safety_decision=SafetyDecision.DENY,
)
assert result.success is False
assert result.error == "Permission denied"
assert result.result is None
def test_tool_result_with_ids(self):
"""Test result with approval and checkpoint IDs."""
result = MCPToolResult(
success=True,
approval_id="approval-123",
checkpoint_id="checkpoint-456",
)
assert result.approval_id == "approval-123"
assert result.checkpoint_id == "checkpoint-456"
def test_tool_result_defaults(self):
"""Test result default values."""
result = MCPToolResult(success=True)
assert result.result is None
assert result.error is None
assert result.safety_decision == SafetyDecision.ALLOW
assert result.execution_time_ms == 0.0
assert result.approval_id is None
assert result.checkpoint_id is None
assert result.metadata == {}
class TestMCPSafetyWrapperClassification:
"""Tests for tool classification."""
def test_classify_file_read(self):
"""Test classifying file read tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("file_read") == ActionType.FILE_READ
assert wrapper._classify_tool("get_file") == ActionType.FILE_READ
assert wrapper._classify_tool("list_files") == ActionType.FILE_READ
assert wrapper._classify_tool("search_file") == ActionType.FILE_READ
def test_classify_file_write(self):
"""Test classifying file write tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("file_write") == ActionType.FILE_WRITE
assert wrapper._classify_tool("create_file") == ActionType.FILE_WRITE
assert wrapper._classify_tool("update_file") == ActionType.FILE_WRITE
def test_classify_file_delete(self):
"""Test classifying file delete tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("file_delete") == ActionType.FILE_DELETE
assert wrapper._classify_tool("remove_file") == ActionType.FILE_DELETE
def test_classify_database_read(self):
"""Test classifying database read tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("database_query") == ActionType.DATABASE_QUERY
assert wrapper._classify_tool("db_read") == ActionType.DATABASE_QUERY
assert wrapper._classify_tool("query_database") == ActionType.DATABASE_QUERY
def test_classify_database_mutate(self):
"""Test classifying database mutate tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("database_write") == ActionType.DATABASE_MUTATE
assert wrapper._classify_tool("db_update") == ActionType.DATABASE_MUTATE
assert wrapper._classify_tool("database_delete") == ActionType.DATABASE_MUTATE
def test_classify_shell_command(self):
"""Test classifying shell command tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("shell_execute") == ActionType.SHELL_COMMAND
assert wrapper._classify_tool("exec_command") == ActionType.SHELL_COMMAND
assert wrapper._classify_tool("bash_run") == ActionType.SHELL_COMMAND
def test_classify_git_operation(self):
"""Test classifying git tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("git_commit") == ActionType.GIT_OPERATION
assert wrapper._classify_tool("git_push") == ActionType.GIT_OPERATION
assert wrapper._classify_tool("git_status") == ActionType.GIT_OPERATION
def test_classify_network_request(self):
"""Test classifying network tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("http_get") == ActionType.NETWORK_REQUEST
assert wrapper._classify_tool("fetch_url") == ActionType.NETWORK_REQUEST
assert wrapper._classify_tool("api_request") == ActionType.NETWORK_REQUEST
def test_classify_llm_call(self):
"""Test classifying LLM tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("llm_generate") == ActionType.LLM_CALL
assert wrapper._classify_tool("ai_complete") == ActionType.LLM_CALL
assert wrapper._classify_tool("claude_chat") == ActionType.LLM_CALL
def test_classify_default(self):
"""Test default classification for unknown tools."""
wrapper = MCPSafetyWrapper()
assert wrapper._classify_tool("unknown_tool") == ActionType.TOOL_CALL
assert wrapper._classify_tool("custom_action") == ActionType.TOOL_CALL
class TestMCPSafetyWrapperToolHandlers:
"""Tests for tool handler registration."""
def test_register_tool_handler(self):
"""Test registering a tool handler."""
wrapper = MCPSafetyWrapper()
def handler(path: str) -> str:
return f"Read: {path}"
wrapper.register_tool_handler("file_read", handler)
assert "file_read" in wrapper._tool_handlers
assert wrapper._tool_handlers["file_read"] is handler
def test_register_multiple_handlers(self):
"""Test registering multiple handlers."""
wrapper = MCPSafetyWrapper()
wrapper.register_tool_handler("tool1", lambda: None)
wrapper.register_tool_handler("tool2", lambda: None)
wrapper.register_tool_handler("tool3", lambda: None)
assert len(wrapper._tool_handlers) == 3
def test_overwrite_handler(self):
"""Test overwriting a handler."""
wrapper = MCPSafetyWrapper()
handler1 = lambda: "first" # noqa: E731
handler2 = lambda: "second" # noqa: E731
wrapper.register_tool_handler("tool", handler1)
wrapper.register_tool_handler("tool", handler2)
assert wrapper._tool_handlers["tool"] is handler2
class TestMCPSafetyWrapperExecution:
"""Tests for tool execution."""
@pytest_asyncio.fixture
async def mock_guardian(self):
"""Create a mock SafetyGuardian."""
guardian = AsyncMock()
guardian.validate = AsyncMock()
return guardian
@pytest_asyncio.fixture
async def mock_emergency(self):
"""Create a mock EmergencyControls."""
emergency = AsyncMock()
emergency.check_allowed = AsyncMock()
return emergency
@pytest.mark.asyncio
async def test_execute_allowed(self, mock_guardian, mock_emergency):
"""Test executing an allowed tool call."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
async def handler(path: str) -> dict:
return {"content": f"Data from {path}"}
wrapper.register_tool_handler("file_read", handler)
call = MCPToolCall(
tool_name="file_read",
arguments={"path": "/test.txt"},
project_id="proj-1",
)
result = await wrapper.execute(call, "agent-1")
assert result.success is True
assert result.result == {"content": "Data from /test.txt"}
assert result.safety_decision == SafetyDecision.ALLOW
@pytest.mark.asyncio
async def test_execute_denied(self, mock_guardian, mock_emergency):
"""Test executing a denied tool call."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.DENY,
reasons=["Permission denied", "Rate limit exceeded"],
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
call = MCPToolCall(
tool_name="file_write",
arguments={"path": "/etc/passwd"},
)
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert "Permission denied" in result.error
assert "Rate limit exceeded" in result.error
assert result.safety_decision == SafetyDecision.DENY
@pytest.mark.asyncio
async def test_execute_requires_approval(self, mock_guardian, mock_emergency):
"""Test executing a tool that requires approval."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.REQUIRE_APPROVAL,
reasons=["Destructive operation requires approval"],
approval_id="approval-123",
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
call = MCPToolCall(
tool_name="file_delete",
arguments={"path": "/important.txt"},
)
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert result.safety_decision == SafetyDecision.REQUIRE_APPROVAL
assert result.approval_id == "approval-123"
assert "requires human approval" in result.error
@pytest.mark.asyncio
async def test_execute_emergency_stop(self, mock_guardian, mock_emergency):
"""Test execution blocked by emergency stop."""
mock_emergency.check_allowed.side_effect = EmergencyStopError(
"Emergency stop active"
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
call = MCPToolCall(
tool_name="file_write",
arguments={"path": "/test.txt"},
project_id="proj-1",
)
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert result.safety_decision == SafetyDecision.DENY
assert result.metadata.get("emergency_stop") is True
@pytest.mark.asyncio
async def test_execute_bypass_safety(self, mock_guardian, mock_emergency):
"""Test executing with safety bypass."""
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
async def handler(data: str) -> str:
return f"Processed: {data}"
wrapper.register_tool_handler("custom_tool", handler)
call = MCPToolCall(
tool_name="custom_tool",
arguments={"data": "test"},
)
result = await wrapper.execute(call, "agent-1", bypass_safety=True)
assert result.success is True
assert result.result == "Processed: test"
# Guardian should not be called when bypassing
mock_guardian.validate.assert_not_called()
@pytest.mark.asyncio
async def test_execute_no_handler(self, mock_guardian, mock_emergency):
"""Test executing a tool with no registered handler."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
call = MCPToolCall(
tool_name="unregistered_tool",
arguments={},
)
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert "No handler registered" in result.error
@pytest.mark.asyncio
async def test_execute_handler_exception(self, mock_guardian, mock_emergency):
"""Test handling exceptions from tool handler."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
async def failing_handler() -> None:
raise ValueError("Handler failed!")
wrapper.register_tool_handler("failing_tool", failing_handler)
call = MCPToolCall(
tool_name="failing_tool",
arguments={},
)
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert "Handler failed!" in result.error
# Decision is still ALLOW because the safety check passed
assert result.safety_decision == SafetyDecision.ALLOW
@pytest.mark.asyncio
async def test_execute_sync_handler(self, mock_guardian, mock_emergency):
"""Test executing a synchronous handler."""
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
def sync_handler(value: int) -> int:
return value * 2
wrapper.register_tool_handler("sync_tool", sync_handler)
call = MCPToolCall(
tool_name="sync_tool",
arguments={"value": 21},
)
result = await wrapper.execute(call, "agent-1")
assert result.success is True
assert result.result == 42
class TestBuildActionRequest:
"""Tests for _build_action_request."""
def test_build_action_request_basic(self):
"""Test building a basic action request."""
wrapper = MCPSafetyWrapper()
call = MCPToolCall(
tool_name="file_read",
arguments={"path": "/test.txt"},
project_id="proj-1",
)
action = wrapper._build_action_request(call, "agent-1", AutonomyLevel.MILESTONE)
assert action.action_type == ActionType.FILE_READ
assert action.tool_name == "file_read"
assert action.arguments == {"path": "/test.txt"}
assert action.resource == "/test.txt"
assert action.metadata.agent_id == "agent-1"
assert action.metadata.project_id == "proj-1"
assert action.metadata.autonomy_level == AutonomyLevel.MILESTONE
def test_build_action_request_with_context(self):
"""Test building action request with session context."""
wrapper = MCPSafetyWrapper()
call = MCPToolCall(
tool_name="database_query",
arguments={"resource": "users", "query": "SELECT *"},
context={"session_id": "sess-123"},
project_id="proj-2",
)
action = wrapper._build_action_request(
call, "agent-2", AutonomyLevel.AUTONOMOUS
)
assert action.resource == "users"
assert action.metadata.session_id == "sess-123"
assert action.metadata.autonomy_level == AutonomyLevel.AUTONOMOUS
def test_build_action_request_no_resource(self):
"""Test building action request without resource."""
wrapper = MCPSafetyWrapper()
call = MCPToolCall(
tool_name="llm_generate",
arguments={"prompt": "Hello"},
)
action = wrapper._build_action_request(
call, "agent-1", AutonomyLevel.FULL_CONTROL
)
assert action.resource is None
class TestElapsedTime:
"""Tests for _elapsed_ms helper."""
def test_elapsed_ms(self):
"""Test calculating elapsed time."""
wrapper = MCPSafetyWrapper()
start = datetime.utcnow() - timedelta(milliseconds=100)
elapsed = wrapper._elapsed_ms(start)
# Should be at least 100ms, but allow some tolerance
assert elapsed >= 99
assert elapsed < 200
class TestSafeToolExecutor:
"""Tests for SafeToolExecutor context manager."""
@pytest.mark.asyncio
async def test_executor_execute(self):
"""Test executing within context manager."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
async def handler() -> str:
return "success"
wrapper.register_tool_handler("test_tool", handler)
call = MCPToolCall(tool_name="test_tool", arguments={})
async with SafeToolExecutor(wrapper, call, "agent-1") as executor:
result = await executor.execute()
assert result.success is True
assert result.result == "success"
@pytest.mark.asyncio
async def test_executor_result_property(self):
"""Test accessing result via property."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
wrapper.register_tool_handler("tool", lambda: "data")
call = MCPToolCall(tool_name="tool", arguments={})
executor = SafeToolExecutor(wrapper, call, "agent-1")
# Before execution
assert executor.result is None
async with executor:
await executor.execute()
# After execution
assert executor.result is not None
assert executor.result.success is True
@pytest.mark.asyncio
async def test_executor_with_autonomy_level(self):
"""Test executor with custom autonomy level."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
wrapper.register_tool_handler("tool", lambda: None)
call = MCPToolCall(tool_name="tool", arguments={})
async with SafeToolExecutor(
wrapper, call, "agent-1", AutonomyLevel.AUTONOMOUS
) as executor:
await executor.execute()
# Check that guardian was called with correct autonomy level
mock_guardian.validate.assert_called_once()
action = mock_guardian.validate.call_args[0][0]
assert action.metadata.autonomy_level == AutonomyLevel.AUTONOMOUS
class TestCreateMCPWrapper:
"""Tests for create_mcp_wrapper factory function."""
@pytest.mark.asyncio
async def test_create_wrapper_with_guardian(self):
"""Test creating wrapper with provided guardian."""
mock_guardian = AsyncMock()
with patch(
"app.services.safety.mcp.integration.get_emergency_controls"
) as mock_get_emergency:
mock_get_emergency.return_value = AsyncMock()
wrapper = await create_mcp_wrapper(guardian=mock_guardian)
assert wrapper._guardian is mock_guardian
@pytest.mark.asyncio
async def test_create_wrapper_default_guardian(self):
"""Test creating wrapper with default guardian."""
with (
patch(
"app.services.safety.mcp.integration.get_safety_guardian"
) as mock_get_guardian,
patch(
"app.services.safety.mcp.integration.get_emergency_controls"
) as mock_get_emergency,
):
mock_guardian = AsyncMock()
mock_get_guardian.return_value = mock_guardian
mock_get_emergency.return_value = AsyncMock()
wrapper = await create_mcp_wrapper()
assert wrapper._guardian is mock_guardian
mock_get_guardian.assert_called_once()
class TestLazyGetters:
"""Tests for lazy getter methods."""
@pytest.mark.asyncio
async def test_get_guardian_lazy(self):
"""Test lazy guardian initialization."""
wrapper = MCPSafetyWrapper()
with patch(
"app.services.safety.mcp.integration.get_safety_guardian"
) as mock_get:
mock_guardian = AsyncMock()
mock_get.return_value = mock_guardian
guardian = await wrapper._get_guardian()
assert guardian is mock_guardian
mock_get.assert_called_once()
@pytest.mark.asyncio
async def test_get_guardian_cached(self):
"""Test guardian is cached after first access."""
mock_guardian = AsyncMock()
wrapper = MCPSafetyWrapper(guardian=mock_guardian)
guardian = await wrapper._get_guardian()
assert guardian is mock_guardian
@pytest.mark.asyncio
async def test_get_emergency_controls_lazy(self):
"""Test lazy emergency controls initialization."""
wrapper = MCPSafetyWrapper()
with patch(
"app.services.safety.mcp.integration.get_emergency_controls"
) as mock_get:
mock_emergency = AsyncMock()
mock_get.return_value = mock_emergency
emergency = await wrapper._get_emergency_controls()
assert emergency is mock_emergency
mock_get.assert_called_once()
@pytest.mark.asyncio
async def test_get_emergency_controls_cached(self):
"""Test emergency controls is cached after first access."""
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(emergency_controls=mock_emergency)
emergency = await wrapper._get_emergency_controls()
assert emergency is mock_emergency
class TestEdgeCases:
"""Tests for edge cases and error handling."""
@pytest.mark.asyncio
async def test_execute_with_safety_error(self):
"""Test handling SafetyError from guardian."""
from app.services.safety.exceptions import SafetyError
mock_guardian = AsyncMock()
mock_guardian.validate.side_effect = SafetyError("Internal safety error")
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
call = MCPToolCall(tool_name="test", arguments={})
result = await wrapper.execute(call, "agent-1")
assert result.success is False
assert "Internal safety error" in result.error
assert result.safety_decision == SafetyDecision.DENY
@pytest.mark.asyncio
async def test_execute_with_checkpoint_id(self):
"""Test that checkpoint_id is propagated to result."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id="checkpoint-abc",
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
wrapper.register_tool_handler("tool", lambda: "result")
call = MCPToolCall(tool_name="tool", arguments={})
result = await wrapper.execute(call, "agent-1")
assert result.success is True
assert result.checkpoint_id == "checkpoint-abc"
def test_destructive_tools_constant(self):
"""Test DESTRUCTIVE_TOOLS class constant."""
assert "file_write" in MCPSafetyWrapper.DESTRUCTIVE_TOOLS
assert "file_delete" in MCPSafetyWrapper.DESTRUCTIVE_TOOLS
assert "shell_execute" in MCPSafetyWrapper.DESTRUCTIVE_TOOLS
assert "git_push" in MCPSafetyWrapper.DESTRUCTIVE_TOOLS
def test_read_only_tools_constant(self):
"""Test READ_ONLY_TOOLS class constant."""
assert "file_read" in MCPSafetyWrapper.READ_ONLY_TOOLS
assert "database_query" in MCPSafetyWrapper.READ_ONLY_TOOLS
assert "git_status" in MCPSafetyWrapper.READ_ONLY_TOOLS
assert "search" in MCPSafetyWrapper.READ_ONLY_TOOLS
@pytest.mark.asyncio
async def test_scope_with_project_id(self):
"""Test that scope is set correctly with project_id."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
wrapper.register_tool_handler("tool", lambda: None)
call = MCPToolCall(
tool_name="tool",
arguments={},
project_id="proj-123",
)
await wrapper.execute(call, "agent-1")
# Verify emergency check was called with project scope
mock_emergency.check_allowed.assert_called_once()
call_kwargs = mock_emergency.check_allowed.call_args
assert "project:proj-123" in str(call_kwargs)
@pytest.mark.asyncio
async def test_scope_without_project_id(self):
"""Test that scope falls back to agent when no project_id."""
mock_guardian = AsyncMock()
mock_guardian.validate.return_value = MagicMock(
decision=SafetyDecision.ALLOW,
reasons=[],
approval_id=None,
checkpoint_id=None,
)
mock_emergency = AsyncMock()
wrapper = MCPSafetyWrapper(
guardian=mock_guardian,
emergency_controls=mock_emergency,
)
wrapper.register_tool_handler("tool", lambda: None)
call = MCPToolCall(
tool_name="tool",
arguments={},
# No project_id
)
await wrapper.execute(call, "agent-555")
# Verify emergency check was called with agent scope
mock_emergency.check_allowed.assert_called_once()
call_kwargs = mock_emergency.check_allowed.call_args
assert "agent:agent-555" in str(call_kwargs)

747
backend/tests/services/safety/test_metrics.py Normal file
View File

@@ -0,0 +1,747 @@
"""
Tests for Safety Metrics Collector.
Tests cover:
- MetricType, MetricValue, HistogramBucket data structures
- SafetyMetrics counters, gauges, histograms
- Prometheus format export
- Summary and reset operations
- Singleton pattern and convenience functions
"""
import pytest
import pytest_asyncio
from app.services.safety.metrics.collector import (
HistogramBucket,
MetricType,
MetricValue,
SafetyMetrics,
get_safety_metrics,
record_mcp_call,
record_validation,
)
class TestMetricType:
"""Tests for MetricType enum."""
def test_metric_types_exist(self):
"""Test all metric types are defined."""
assert MetricType.COUNTER == "counter"
assert MetricType.GAUGE == "gauge"
assert MetricType.HISTOGRAM == "histogram"
def test_metric_type_is_string(self):
"""Test MetricType values are strings."""
assert isinstance(MetricType.COUNTER.value, str)
assert isinstance(MetricType.GAUGE.value, str)
assert isinstance(MetricType.HISTOGRAM.value, str)
class TestMetricValue:
"""Tests for MetricValue dataclass."""
def test_metric_value_creation(self):
"""Test creating a metric value."""
mv = MetricValue(
name="test_metric",
metric_type=MetricType.COUNTER,
value=42.0,
labels={"env": "test"},
)
assert mv.name == "test_metric"
assert mv.metric_type == MetricType.COUNTER
assert mv.value == 42.0
assert mv.labels == {"env": "test"}
assert mv.timestamp is not None
def test_metric_value_defaults(self):
"""Test metric value default values."""
mv = MetricValue(
name="test",
metric_type=MetricType.GAUGE,
value=0.0,
)
assert mv.labels == {}
assert mv.timestamp is not None
class TestHistogramBucket:
"""Tests for HistogramBucket dataclass."""
def test_histogram_bucket_creation(self):
"""Test creating a histogram bucket."""
bucket = HistogramBucket(le=0.5, count=10)
assert bucket.le == 0.5
assert bucket.count == 10
def test_histogram_bucket_defaults(self):
"""Test histogram bucket default count."""
bucket = HistogramBucket(le=1.0)
assert bucket.le == 1.0
assert bucket.count == 0
def test_histogram_bucket_infinity(self):
"""Test histogram bucket with infinity."""
bucket = HistogramBucket(le=float("inf"))
assert bucket.le == float("inf")
class TestSafetyMetricsCounters:
"""Tests for SafetyMetrics counter methods."""
@pytest_asyncio.fixture
async def metrics(self):
"""Create fresh metrics instance."""
return SafetyMetrics()
@pytest.mark.asyncio
async def test_inc_validations(self, metrics):
"""Test incrementing validations counter."""
await metrics.inc_validations("allow")
await metrics.inc_validations("allow")
await metrics.inc_validations("deny", agent_id="agent-1")
summary = await metrics.get_summary()
assert summary["total_validations"] == 3
assert summary["denied_validations"] == 1
@pytest.mark.asyncio
async def test_inc_approvals_requested(self, metrics):
"""Test incrementing approval requests counter."""
await metrics.inc_approvals_requested("normal")
await metrics.inc_approvals_requested("urgent")
await metrics.inc_approvals_requested() # default
summary = await metrics.get_summary()
assert summary["approval_requests"] == 3
@pytest.mark.asyncio
async def test_inc_approvals_granted(self, metrics):
"""Test incrementing approvals granted counter."""
await metrics.inc_approvals_granted()
await metrics.inc_approvals_granted()
summary = await metrics.get_summary()
assert summary["approvals_granted"] == 2
@pytest.mark.asyncio
async def test_inc_approvals_denied(self, metrics):
"""Test incrementing approvals denied counter."""
await metrics.inc_approvals_denied("timeout")
await metrics.inc_approvals_denied("policy")
await metrics.inc_approvals_denied() # default manual
summary = await metrics.get_summary()
assert summary["approvals_denied"] == 3
@pytest.mark.asyncio
async def test_inc_rate_limit_exceeded(self, metrics):
"""Test incrementing rate limit exceeded counter."""
await metrics.inc_rate_limit_exceeded("requests_per_minute")
await metrics.inc_rate_limit_exceeded("tokens_per_hour")
summary = await metrics.get_summary()
assert summary["rate_limit_hits"] == 2
@pytest.mark.asyncio
async def test_inc_budget_exceeded(self, metrics):
"""Test incrementing budget exceeded counter."""
await metrics.inc_budget_exceeded("daily_cost")
await metrics.inc_budget_exceeded("monthly_tokens")
summary = await metrics.get_summary()
assert summary["budget_exceeded"] == 2
@pytest.mark.asyncio
async def test_inc_loops_detected(self, metrics):
"""Test incrementing loops detected counter."""
await metrics.inc_loops_detected("repetition")
await metrics.inc_loops_detected("pattern")
summary = await metrics.get_summary()
assert summary["loops_detected"] == 2
@pytest.mark.asyncio
async def test_inc_emergency_events(self, metrics):
"""Test incrementing emergency events counter."""
await metrics.inc_emergency_events("pause", "project-1")
await metrics.inc_emergency_events("stop", "agent-2")
summary = await metrics.get_summary()
assert summary["emergency_events"] == 2
@pytest.mark.asyncio
async def test_inc_content_filtered(self, metrics):
"""Test incrementing content filtered counter."""
await metrics.inc_content_filtered("profanity", "blocked")
await metrics.inc_content_filtered("pii", "redacted")
summary = await metrics.get_summary()
assert summary["content_filtered"] == 2
@pytest.mark.asyncio
async def test_inc_checkpoints_created(self, metrics):
"""Test incrementing checkpoints created counter."""
await metrics.inc_checkpoints_created()
await metrics.inc_checkpoints_created()
await metrics.inc_checkpoints_created()
summary = await metrics.get_summary()
assert summary["checkpoints_created"] == 3
@pytest.mark.asyncio
async def test_inc_rollbacks_executed(self, metrics):
"""Test incrementing rollbacks executed counter."""
await metrics.inc_rollbacks_executed(success=True)
await metrics.inc_rollbacks_executed(success=False)
summary = await metrics.get_summary()
assert summary["rollbacks_executed"] == 2
@pytest.mark.asyncio
async def test_inc_mcp_calls(self, metrics):
"""Test incrementing MCP calls counter."""
await metrics.inc_mcp_calls("search_knowledge", success=True)
await metrics.inc_mcp_calls("run_code", success=False)
summary = await metrics.get_summary()
assert summary["mcp_calls"] == 2
class TestSafetyMetricsGauges:
"""Tests for SafetyMetrics gauge methods."""
@pytest_asyncio.fixture
async def metrics(self):
"""Create fresh metrics instance."""
return SafetyMetrics()
@pytest.mark.asyncio
async def test_set_budget_remaining(self, metrics):
"""Test setting budget remaining gauge."""
await metrics.set_budget_remaining("project-1", "daily_cost", 50.0)
all_metrics = await metrics.get_all_metrics()
gauge_metrics = [m for m in all_metrics if m.name == "safety_budget_remaining"]
assert len(gauge_metrics) == 1
assert gauge_metrics[0].value == 50.0
assert gauge_metrics[0].labels["scope"] == "project-1"
assert gauge_metrics[0].labels["budget_type"] == "daily_cost"
@pytest.mark.asyncio
async def test_set_rate_limit_remaining(self, metrics):
"""Test setting rate limit remaining gauge."""
await metrics.set_rate_limit_remaining("agent-1", "requests_per_minute", 45)
all_metrics = await metrics.get_all_metrics()
gauge_metrics = [
m for m in all_metrics if m.name == "safety_rate_limit_remaining"
]
assert len(gauge_metrics) == 1
assert gauge_metrics[0].value == 45.0
@pytest.mark.asyncio
async def test_set_pending_approvals(self, metrics):
"""Test setting pending approvals gauge."""
await metrics.set_pending_approvals(5)
summary = await metrics.get_summary()
assert summary["pending_approvals"] == 5
@pytest.mark.asyncio
async def test_set_active_checkpoints(self, metrics):
"""Test setting active checkpoints gauge."""
await metrics.set_active_checkpoints(3)
summary = await metrics.get_summary()
assert summary["active_checkpoints"] == 3
@pytest.mark.asyncio
async def test_set_emergency_state(self, metrics):
"""Test setting emergency state gauge."""
await metrics.set_emergency_state("project-1", "normal")
await metrics.set_emergency_state("project-2", "paused")
await metrics.set_emergency_state("project-3", "stopped")
await metrics.set_emergency_state("project-4", "unknown")
all_metrics = await metrics.get_all_metrics()
state_metrics = [m for m in all_metrics if m.name == "safety_emergency_state"]
assert len(state_metrics) == 4
# Check state values
values_by_scope = {m.labels["scope"]: m.value for m in state_metrics}
assert values_by_scope["project-1"] == 0.0 # normal
assert values_by_scope["project-2"] == 1.0 # paused
assert values_by_scope["project-3"] == 2.0 # stopped
assert values_by_scope["project-4"] == -1.0 # unknown
class TestSafetyMetricsHistograms:
"""Tests for SafetyMetrics histogram methods."""
@pytest_asyncio.fixture
async def metrics(self):
"""Create fresh metrics instance."""
return SafetyMetrics()
@pytest.mark.asyncio
async def test_observe_validation_latency(self, metrics):
"""Test observing validation latency."""
await metrics.observe_validation_latency(0.05)
await metrics.observe_validation_latency(0.15)
await metrics.observe_validation_latency(0.5)
all_metrics = await metrics.get_all_metrics()
count_metric = next(
(m for m in all_metrics if m.name == "validation_latency_seconds_count"),
None,
)
assert count_metric is not None
assert count_metric.value == 3.0
sum_metric = next(
(m for m in all_metrics if m.name == "validation_latency_seconds_sum"),
None,
)
assert sum_metric is not None
assert abs(sum_metric.value - 0.7) < 0.001
@pytest.mark.asyncio
async def test_observe_approval_latency(self, metrics):
"""Test observing approval latency."""
await metrics.observe_approval_latency(1.5)
await metrics.observe_approval_latency(3.0)
all_metrics = await metrics.get_all_metrics()
count_metric = next(
(m for m in all_metrics if m.name == "approval_latency_seconds_count"),
None,
)
assert count_metric is not None
assert count_metric.value == 2.0
@pytest.mark.asyncio
async def test_observe_mcp_execution_latency(self, metrics):
"""Test observing MCP execution latency."""
await metrics.observe_mcp_execution_latency(0.02)
all_metrics = await metrics.get_all_metrics()
count_metric = next(
(m for m in all_metrics if m.name == "mcp_execution_latency_seconds_count"),
None,
)
assert count_metric is not None
assert count_metric.value == 1.0
@pytest.mark.asyncio
async def test_histogram_bucket_updates(self, metrics):
"""Test that histogram buckets are updated correctly."""
# Add values to test bucket distribution
await metrics.observe_validation_latency(0.005) # <= 0.01
await metrics.observe_validation_latency(0.03) # <= 0.05
await metrics.observe_validation_latency(0.07) # <= 0.1
await metrics.observe_validation_latency(15.0) # <= inf
prometheus = await metrics.get_prometheus_format()
# Check that bucket counts are in output
assert "validation_latency_seconds_bucket" in prometheus
assert "le=" in prometheus
class TestSafetyMetricsExport:
"""Tests for SafetyMetrics export methods."""
@pytest_asyncio.fixture
async def metrics(self):
"""Create fresh metrics instance with some data."""
m = SafetyMetrics()
# Add some counters
await m.inc_validations("allow")
await m.inc_validations("deny", agent_id="agent-1")
# Add some gauges
await m.set_pending_approvals(3)
await m.set_budget_remaining("proj-1", "daily", 100.0)
# Add some histogram values
await m.observe_validation_latency(0.1)
return m
@pytest.mark.asyncio
async def test_get_all_metrics(self, metrics):
"""Test getting all metrics."""
all_metrics = await metrics.get_all_metrics()
assert len(all_metrics) > 0
assert all(isinstance(m, MetricValue) for m in all_metrics)
# Check we have different types
types = {m.metric_type for m in all_metrics}
assert MetricType.COUNTER in types
assert MetricType.GAUGE in types
@pytest.mark.asyncio
async def test_get_prometheus_format(self, metrics):
"""Test Prometheus format export."""
output = await metrics.get_prometheus_format()
assert isinstance(output, str)
assert "# TYPE" in output
assert "counter" in output
assert "gauge" in output
assert "safety_validations_total" in output
assert "safety_pending_approvals" in output
@pytest.mark.asyncio
async def test_prometheus_format_with_labels(self, metrics):
"""Test Prometheus format includes labels correctly."""
output = await metrics.get_prometheus_format()
# Counter with labels
assert "decision=allow" in output or "decision=deny" in output
@pytest.mark.asyncio
async def test_prometheus_format_histogram_buckets(self, metrics):
"""Test Prometheus format includes histogram buckets."""
output = await metrics.get_prometheus_format()
assert "histogram" in output
assert "_bucket" in output
assert "le=" in output
assert "+Inf" in output
@pytest.mark.asyncio
async def test_get_summary(self, metrics):
"""Test getting summary."""
summary = await metrics.get_summary()
assert "total_validations" in summary
assert "denied_validations" in summary
assert "approval_requests" in summary
assert "pending_approvals" in summary
assert "active_checkpoints" in summary
assert summary["total_validations"] == 2
assert summary["denied_validations"] == 1
assert summary["pending_approvals"] == 3
@pytest.mark.asyncio
async def test_summary_empty_counters(self):
"""Test summary with no data."""
metrics = SafetyMetrics()
summary = await metrics.get_summary()
assert summary["total_validations"] == 0
assert summary["denied_validations"] == 0
assert summary["pending_approvals"] == 0
class TestSafetyMetricsReset:
"""Tests for SafetyMetrics reset."""
@pytest.mark.asyncio
async def test_reset_clears_counters(self):
"""Test reset clears all counters."""
metrics = SafetyMetrics()
await metrics.inc_validations("allow")
await metrics.inc_approvals_granted()
await metrics.set_pending_approvals(5)
await metrics.observe_validation_latency(0.1)
await metrics.reset()
summary = await metrics.get_summary()
assert summary["total_validations"] == 0
assert summary["approvals_granted"] == 0
assert summary["pending_approvals"] == 0
@pytest.mark.asyncio
async def test_reset_reinitializes_histogram_buckets(self):
"""Test reset reinitializes histogram buckets."""
metrics = SafetyMetrics()
await metrics.observe_validation_latency(0.1)
await metrics.reset()
# After reset, histogram buckets should be reinitialized
prometheus = await metrics.get_prometheus_format()
assert "validation_latency_seconds" in prometheus
class TestParseLabels:
"""Tests for _parse_labels helper method."""
def test_parse_empty_labels(self):
"""Test parsing empty labels string."""
metrics = SafetyMetrics()
result = metrics._parse_labels("")
assert result == {}
def test_parse_single_label(self):
"""Test parsing single label."""
metrics = SafetyMetrics()
result = metrics._parse_labels("key=value")
assert result == {"key": "value"}
def test_parse_multiple_labels(self):
"""Test parsing multiple labels."""
metrics = SafetyMetrics()
result = metrics._parse_labels("a=1,b=2,c=3")
assert result == {"a": "1", "b": "2", "c": "3"}
def test_parse_labels_with_spaces(self):
"""Test parsing labels with spaces."""
metrics = SafetyMetrics()
result = metrics._parse_labels(" key = value , foo = bar ")
assert result == {"key": "value", "foo": "bar"}
def test_parse_labels_with_equals_in_value(self):
"""Test parsing labels with = in value."""
metrics = SafetyMetrics()
result = metrics._parse_labels("query=a=b")
assert result == {"query": "a=b"}
def test_parse_invalid_label(self):
"""Test parsing invalid label without equals."""
metrics = SafetyMetrics()
result = metrics._parse_labels("no_equals")
assert result == {}
class TestHistogramBucketInit:
"""Tests for histogram bucket initialization."""
def test_histogram_buckets_initialized(self):
"""Test that histogram buckets are initialized."""
metrics = SafetyMetrics()
assert "validation_latency_seconds" in metrics._histogram_buckets
assert "approval_latency_seconds" in metrics._histogram_buckets
assert "mcp_execution_latency_seconds" in metrics._histogram_buckets
def test_histogram_buckets_have_correct_values(self):
"""Test histogram buckets have correct boundary values."""
metrics = SafetyMetrics()
buckets = metrics._histogram_buckets["validation_latency_seconds"]
# Check first few and last bucket
assert buckets[0].le == 0.01
assert buckets[1].le == 0.05
assert buckets[-1].le == float("inf")
# Check all have zero initial count
assert all(b.count == 0 for b in buckets)
class TestSingletonAndConvenience:
"""Tests for singleton pattern and convenience functions."""
@pytest.mark.asyncio
async def test_get_safety_metrics_returns_same_instance(self):
"""Test get_safety_metrics returns singleton."""
# Reset the module-level singleton for this test
import app.services.safety.metrics.collector as collector_module
collector_module._metrics = None
m1 = await get_safety_metrics()
m2 = await get_safety_metrics()
assert m1 is m2
@pytest.mark.asyncio
async def test_record_validation_convenience(self):
"""Test record_validation convenience function."""
import app.services.safety.metrics.collector as collector_module
collector_module._metrics = None # Reset
await record_validation("allow")
await record_validation("deny", agent_id="test-agent")
metrics = await get_safety_metrics()
summary = await metrics.get_summary()
assert summary["total_validations"] == 2
assert summary["denied_validations"] == 1
@pytest.mark.asyncio
async def test_record_mcp_call_convenience(self):
"""Test record_mcp_call convenience function."""
import app.services.safety.metrics.collector as collector_module
collector_module._metrics = None # Reset
await record_mcp_call("search_knowledge", success=True, latency_ms=50)
await record_mcp_call("run_code", success=False, latency_ms=100)
metrics = await get_safety_metrics()
summary = await metrics.get_summary()
assert summary["mcp_calls"] == 2
class TestConcurrency:
"""Tests for concurrent metric updates."""
@pytest.mark.asyncio
async def test_concurrent_counter_increments(self):
"""Test concurrent counter increments are safe."""
import asyncio
metrics = SafetyMetrics()
async def increment_many():
for _ in range(100):
await metrics.inc_validations("allow")
# Run 10 concurrent tasks each incrementing 100 times
await asyncio.gather(*[increment_many() for _ in range(10)])
summary = await metrics.get_summary()
assert summary["total_validations"] == 1000
@pytest.mark.asyncio
async def test_concurrent_gauge_updates(self):
"""Test concurrent gauge updates are safe."""
import asyncio
metrics = SafetyMetrics()
async def update_gauge(value):
await metrics.set_pending_approvals(value)
# Run concurrent gauge updates
await asyncio.gather(*[update_gauge(i) for i in range(100)])
# Final value should be one of the updates (last one wins)
summary = await metrics.get_summary()
assert 0 <= summary["pending_approvals"] < 100
@pytest.mark.asyncio
async def test_concurrent_histogram_observations(self):
"""Test concurrent histogram observations are safe."""
import asyncio
metrics = SafetyMetrics()
async def observe_many():
for i in range(100):
await metrics.observe_validation_latency(i / 1000)
await asyncio.gather(*[observe_many() for _ in range(10)])
all_metrics = await metrics.get_all_metrics()
count_metric = next(
(m for m in all_metrics if m.name == "validation_latency_seconds_count"),
None,
)
assert count_metric is not None
assert count_metric.value == 1000.0
class TestEdgeCases:
"""Tests for edge cases."""
@pytest.mark.asyncio
async def test_very_large_counter_value(self):
"""Test handling very large counter values."""
metrics = SafetyMetrics()
for _ in range(10000):
await metrics.inc_validations("allow")
summary = await metrics.get_summary()
assert summary["total_validations"] == 10000
@pytest.mark.asyncio
async def test_zero_and_negative_gauge_values(self):
"""Test zero and negative gauge values."""
metrics = SafetyMetrics()
await metrics.set_budget_remaining("project", "cost", 0.0)
await metrics.set_budget_remaining("project2", "cost", -10.0)
all_metrics = await metrics.get_all_metrics()
gauges = [m for m in all_metrics if m.name == "safety_budget_remaining"]
values = {m.labels.get("scope"): m.value for m in gauges}
assert values["project"] == 0.0
assert values["project2"] == -10.0
@pytest.mark.asyncio
async def test_very_small_histogram_values(self):
"""Test very small histogram values."""
metrics = SafetyMetrics()
await metrics.observe_validation_latency(0.0001) # 0.1ms
all_metrics = await metrics.get_all_metrics()
sum_metric = next(
(m for m in all_metrics if m.name == "validation_latency_seconds_sum"),
None,
)
assert sum_metric is not None
assert abs(sum_metric.value - 0.0001) < 0.00001
@pytest.mark.asyncio
async def test_special_characters_in_labels(self):
"""Test special characters in label values."""
metrics = SafetyMetrics()
await metrics.inc_validations("allow", agent_id="agent/with/slashes")
all_metrics = await metrics.get_all_metrics()
counters = [m for m in all_metrics if m.name == "safety_validations_total"]
# Should have the metric with special chars
assert len(counters) > 0
@pytest.mark.asyncio
async def test_empty_histogram_export(self):
"""Test exporting histogram with no observations."""
metrics = SafetyMetrics()
# No observations, but histogram buckets should still exist
prometheus = await metrics.get_prometheus_format()
assert "validation_latency_seconds" in prometheus
assert "le=" in prometheus
@pytest.mark.asyncio
async def test_prometheus_format_empty_label_value(self):
"""Test Prometheus format with empty label metrics."""
metrics = SafetyMetrics()
await metrics.inc_approvals_granted() # Uses empty string as label
prometheus = await metrics.get_prometheus_format()
assert "safety_approvals_granted_total" in prometheus
@pytest.mark.asyncio
async def test_multiple_resets(self):
"""Test multiple resets don't cause issues."""
metrics = SafetyMetrics()
await metrics.inc_validations("allow")
await metrics.reset()
await metrics.reset()
await metrics.reset()
summary = await metrics.get_summary()
assert summary["total_validations"] == 0

933
backend/tests/services/safety/test_permissions.py Normal file
View File

@@ -0,0 +1,933 @@
"""Tests for Permission Manager.
Tests cover:
- PermissionGrant: creation, expiry, matching, hierarchy
- PermissionManager: grant, revoke, check, require, list, defaults
- Edge cases: wildcards, expiration, default deny/allow
"""
from datetime import datetime, timedelta
import pytest
import pytest_asyncio
from app.services.safety.exceptions import PermissionDeniedError
from app.services.safety.models import (
ActionMetadata,
ActionRequest,
ActionType,
PermissionLevel,
ResourceType,
)
from app.services.safety.permissions.manager import PermissionGrant, PermissionManager
# ============================================================================
# Fixtures
# ============================================================================
@pytest.fixture
def action_metadata() -> ActionMetadata:
"""Create standard action metadata for tests."""
return ActionMetadata(
agent_id="test-agent",
project_id="test-project",
session_id="test-session",
)
@pytest_asyncio.fixture
async def permission_manager() -> PermissionManager:
"""Create a PermissionManager for testing."""
return PermissionManager(default_deny=True)
@pytest_asyncio.fixture
async def permissive_manager() -> PermissionManager:
"""Create a PermissionManager with default_deny=False."""
return PermissionManager(default_deny=False)
# ============================================================================
# PermissionGrant Tests
# ============================================================================
class TestPermissionGrant:
"""Tests for the PermissionGrant class."""
def test_grant_creation(self) -> None:
"""Test basic grant creation."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
granted_by="admin",
reason="Read access to data directory",
)
assert grant.id is not None
assert grant.agent_id == "agent-1"
assert grant.resource_pattern == "/data/*"
assert grant.resource_type == ResourceType.FILE
assert grant.level == PermissionLevel.READ
assert grant.granted_by == "admin"
assert grant.reason == "Read access to data directory"
assert grant.expires_at is None
assert grant.created_at is not None
def test_grant_with_expiration(self) -> None:
"""Test grant with expiration time."""
future = datetime.utcnow() + timedelta(hours=1)
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
expires_at=future,
)
assert grant.expires_at == future
assert grant.is_expired() is False
def test_is_expired_no_expiration(self) -> None:
"""Test is_expired with no expiration set."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
assert grant.is_expired() is False
def test_is_expired_future(self) -> None:
"""Test is_expired with future expiration."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
expires_at=datetime.utcnow() + timedelta(hours=1),
)
assert grant.is_expired() is False
def test_is_expired_past(self) -> None:
"""Test is_expired with past expiration."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
expires_at=datetime.utcnow() - timedelta(hours=1),
)
assert grant.is_expired() is True
def test_matches_exact(self) -> None:
"""Test matching with exact pattern."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="/data/file.txt",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
assert grant.matches("/data/file.txt", ResourceType.FILE) is True
assert grant.matches("/data/other.txt", ResourceType.FILE) is False
def test_matches_wildcard(self) -> None:
"""Test matching with wildcard pattern."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
assert grant.matches("/data/file.txt", ResourceType.FILE) is True
# fnmatch's * matches everything including /
assert grant.matches("/data/subdir/file.txt", ResourceType.FILE) is True
assert grant.matches("/other/file.txt", ResourceType.FILE) is False
def test_matches_recursive_wildcard(self) -> None:
"""Test matching with recursive pattern."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="/data/**",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# fnmatch treats ** similar to * - both match everything including /
assert grant.matches("/data/file.txt", ResourceType.FILE) is True
assert grant.matches("/data/subdir/file.txt", ResourceType.FILE) is True
def test_matches_wrong_resource_type(self) -> None:
"""Test matching fails with wrong resource type."""
grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Same pattern but different resource type
assert grant.matches("/data/table", ResourceType.DATABASE) is False
def test_allows_hierarchy(self) -> None:
"""Test permission level hierarchy."""
admin_grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.ADMIN,
)
# ADMIN allows all levels
assert admin_grant.allows(PermissionLevel.NONE) is True
assert admin_grant.allows(PermissionLevel.READ) is True
assert admin_grant.allows(PermissionLevel.WRITE) is True
assert admin_grant.allows(PermissionLevel.EXECUTE) is True
assert admin_grant.allows(PermissionLevel.DELETE) is True
assert admin_grant.allows(PermissionLevel.ADMIN) is True
def test_allows_read_only(self) -> None:
"""Test READ grant only allows READ and NONE."""
read_grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
assert read_grant.allows(PermissionLevel.NONE) is True
assert read_grant.allows(PermissionLevel.READ) is True
assert read_grant.allows(PermissionLevel.WRITE) is False
assert read_grant.allows(PermissionLevel.EXECUTE) is False
assert read_grant.allows(PermissionLevel.DELETE) is False
assert read_grant.allows(PermissionLevel.ADMIN) is False
def test_allows_write_includes_read(self) -> None:
"""Test WRITE grant includes READ."""
write_grant = PermissionGrant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.WRITE,
)
assert write_grant.allows(PermissionLevel.READ) is True
assert write_grant.allows(PermissionLevel.WRITE) is True
assert write_grant.allows(PermissionLevel.EXECUTE) is False
# ============================================================================
# PermissionManager Tests
# ============================================================================
class TestPermissionManager:
"""Tests for the PermissionManager class."""
@pytest.mark.asyncio
async def test_grant_creates_permission(
self,
permission_manager: PermissionManager,
) -> None:
"""Test granting a permission."""
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
granted_by="admin",
reason="Read access",
)
assert grant.id is not None
assert grant.agent_id == "agent-1"
assert grant.resource_pattern == "/data/*"
@pytest.mark.asyncio
async def test_grant_with_duration(
self,
permission_manager: PermissionManager,
) -> None:
"""Test granting a temporary permission."""
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
duration_seconds=3600, # 1 hour
)
assert grant.expires_at is not None
assert grant.is_expired() is False
@pytest.mark.asyncio
async def test_revoke_by_id(
self,
permission_manager: PermissionManager,
) -> None:
"""Test revoking a grant by ID."""
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
success = await permission_manager.revoke(grant.id)
assert success is True
# Verify grant is removed
grants = await permission_manager.list_grants(agent_id="agent-1")
assert len(grants) == 0
@pytest.mark.asyncio
async def test_revoke_nonexistent(
self,
permission_manager: PermissionManager,
) -> None:
"""Test revoking a non-existent grant."""
success = await permission_manager.revoke("nonexistent-id")
assert success is False
@pytest.mark.asyncio
async def test_revoke_all_for_agent(
self,
permission_manager: PermissionManager,
) -> None:
"""Test revoking all permissions for an agent."""
# Grant multiple permissions
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/api/*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
)
await permission_manager.grant(
agent_id="agent-2",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
revoked = await permission_manager.revoke_all("agent-1")
assert revoked == 2
# Verify agent-1 grants are gone
grants = await permission_manager.list_grants(agent_id="agent-1")
assert len(grants) == 0
# Verify agent-2 grant remains
grants = await permission_manager.list_grants(agent_id="agent-2")
assert len(grants) == 1
@pytest.mark.asyncio
async def test_revoke_all_no_grants(
self,
permission_manager: PermissionManager,
) -> None:
"""Test revoking all when no grants exist."""
revoked = await permission_manager.revoke_all("nonexistent-agent")
assert revoked == 0
@pytest.mark.asyncio
async def test_check_granted(
self,
permission_manager: PermissionManager,
) -> None:
"""Test checking a granted permission."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert allowed is True
@pytest.mark.asyncio
async def test_check_denied_default_deny(
self,
permission_manager: PermissionManager,
) -> None:
"""Test checking denied with default_deny=True."""
# No grants, should be denied
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert allowed is False
@pytest.mark.asyncio
async def test_check_uses_default_permissions(
self,
permissive_manager: PermissionManager,
) -> None:
"""Test that default permissions apply when default_deny=False."""
# No explicit grants, but FILE default is READ
allowed = await permissive_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert allowed is True
# But WRITE should fail
allowed = await permissive_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.WRITE,
)
assert allowed is False
@pytest.mark.asyncio
async def test_check_shell_denied_by_default(
self,
permissive_manager: PermissionManager,
) -> None:
"""Test SHELL is denied by default (NONE level)."""
allowed = await permissive_manager.check(
agent_id="agent-1",
resource="rm -rf /",
resource_type=ResourceType.SHELL,
required_level=PermissionLevel.EXECUTE,
)
assert allowed is False
@pytest.mark.asyncio
async def test_check_expired_grant_ignored(
self,
permission_manager: PermissionManager,
) -> None:
"""Test that expired grants are ignored in checks."""
# Create an already-expired grant
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
duration_seconds=1, # Very short
)
# Manually expire it
grant.expires_at = datetime.utcnow() - timedelta(seconds=10)
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert allowed is False
@pytest.mark.asyncio
async def test_check_insufficient_level(
self,
permission_manager: PermissionManager,
) -> None:
"""Test check fails when grant level is insufficient."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Try to get WRITE access with only READ grant
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.WRITE,
)
assert allowed is False
@pytest.mark.asyncio
async def test_check_action_file_read(
self,
permission_manager: PermissionManager,
action_metadata: ActionMetadata,
) -> None:
"""Test check_action for file read."""
await permission_manager.grant(
agent_id="test-agent",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
action = ActionRequest(
action_type=ActionType.FILE_READ,
resource="/data/file.txt",
metadata=action_metadata,
)
allowed = await permission_manager.check_action(action)
assert allowed is True
@pytest.mark.asyncio
async def test_check_action_file_write(
self,
permission_manager: PermissionManager,
action_metadata: ActionMetadata,
) -> None:
"""Test check_action for file write."""
await permission_manager.grant(
agent_id="test-agent",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.WRITE,
)
action = ActionRequest(
action_type=ActionType.FILE_WRITE,
resource="/data/file.txt",
metadata=action_metadata,
)
allowed = await permission_manager.check_action(action)
assert allowed is True
@pytest.mark.asyncio
async def test_check_action_uses_tool_name_as_resource(
self,
permission_manager: PermissionManager,
action_metadata: ActionMetadata,
) -> None:
"""Test check_action uses tool_name when resource is None."""
await permission_manager.grant(
agent_id="test-agent",
resource_pattern="search_*",
resource_type=ResourceType.CUSTOM,
level=PermissionLevel.EXECUTE,
)
action = ActionRequest(
action_type=ActionType.TOOL_CALL,
tool_name="search_documents",
resource=None,
metadata=action_metadata,
)
allowed = await permission_manager.check_action(action)
assert allowed is True
@pytest.mark.asyncio
async def test_require_permission_granted(
self,
permission_manager: PermissionManager,
) -> None:
"""Test require_permission doesn't raise when granted."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Should not raise
await permission_manager.require_permission(
agent_id="agent-1",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
@pytest.mark.asyncio
async def test_require_permission_denied(
self,
permission_manager: PermissionManager,
) -> None:
"""Test require_permission raises when denied."""
with pytest.raises(PermissionDeniedError) as exc_info:
await permission_manager.require_permission(
agent_id="agent-1",
resource="/secret/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert "/secret/file.txt" in str(exc_info.value)
assert exc_info.value.agent_id == "agent-1"
assert exc_info.value.required_permission == "read"
@pytest.mark.asyncio
async def test_list_grants_all(
self,
permission_manager: PermissionManager,
) -> None:
"""Test listing all grants."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
await permission_manager.grant(
agent_id="agent-2",
resource_pattern="/api/*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
)
grants = await permission_manager.list_grants()
assert len(grants) == 2
@pytest.mark.asyncio
async def test_list_grants_by_agent(
self,
permission_manager: PermissionManager,
) -> None:
"""Test listing grants filtered by agent."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
await permission_manager.grant(
agent_id="agent-2",
resource_pattern="/api/*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
)
grants = await permission_manager.list_grants(agent_id="agent-1")
assert len(grants) == 1
assert grants[0].agent_id == "agent-1"
@pytest.mark.asyncio
async def test_list_grants_by_resource_type(
self,
permission_manager: PermissionManager,
) -> None:
"""Test listing grants filtered by resource type."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/api/*",
resource_type=ResourceType.API,
level=PermissionLevel.EXECUTE,
)
grants = await permission_manager.list_grants(resource_type=ResourceType.FILE)
assert len(grants) == 1
assert grants[0].resource_type == ResourceType.FILE
@pytest.mark.asyncio
async def test_list_grants_excludes_expired(
self,
permission_manager: PermissionManager,
) -> None:
"""Test that list_grants excludes expired grants."""
# Create expired grant
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/old/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
duration_seconds=1,
)
grant.expires_at = datetime.utcnow() - timedelta(seconds=10)
# Create valid grant
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/new/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
grants = await permission_manager.list_grants()
assert len(grants) == 1
assert grants[0].resource_pattern == "/new/*"
def test_set_default_permission(
self,
) -> None:
"""Test setting default permission level."""
manager = PermissionManager(default_deny=False)
# Default for SHELL is NONE
assert manager._default_permissions[ResourceType.SHELL] == PermissionLevel.NONE
# Change it
manager.set_default_permission(ResourceType.SHELL, PermissionLevel.EXECUTE)
assert (
manager._default_permissions[ResourceType.SHELL] == PermissionLevel.EXECUTE
)
@pytest.mark.asyncio
async def test_set_default_permission_affects_checks(
self,
permissive_manager: PermissionManager,
) -> None:
"""Test that changing default permissions affects checks."""
# Initially SHELL is NONE
allowed = await permissive_manager.check(
agent_id="agent-1",
resource="ls",
resource_type=ResourceType.SHELL,
required_level=PermissionLevel.EXECUTE,
)
assert allowed is False
# Change default
permissive_manager.set_default_permission(
ResourceType.SHELL, PermissionLevel.EXECUTE
)
# Now should be allowed
allowed = await permissive_manager.check(
agent_id="agent-1",
resource="ls",
resource_type=ResourceType.SHELL,
required_level=PermissionLevel.EXECUTE,
)
assert allowed is True
# ============================================================================
# Edge Cases
# ============================================================================
class TestPermissionEdgeCases:
"""Edge cases that could reveal hidden bugs."""
@pytest.mark.asyncio
async def test_multiple_matching_grants(
self,
permission_manager: PermissionManager,
) -> None:
"""Test when multiple grants match - first sufficient one wins."""
# Grant READ on all files
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Also grant WRITE on specific path
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/writable/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.WRITE,
)
# Write on writable path should work
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/data/writable/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.WRITE,
)
assert allowed is True
@pytest.mark.asyncio
async def test_wildcard_all_pattern(
self,
permission_manager: PermissionManager,
) -> None:
"""Test * pattern matches everything."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="*",
resource_type=ResourceType.FILE,
level=PermissionLevel.ADMIN,
)
allowed = await permission_manager.check(
agent_id="agent-1",
resource="/any/path/anywhere/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.DELETE,
)
# fnmatch's * matches everything including /
assert allowed is True
@pytest.mark.asyncio
async def test_question_mark_wildcard(
self,
permission_manager: PermissionManager,
) -> None:
"""Test ? wildcard matches single character."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="file?.txt",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
assert (
await permission_manager.check(
agent_id="agent-1",
resource="file1.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
is True
)
assert (
await permission_manager.check(
agent_id="agent-1",
resource="file10.txt", # Two characters, won't match
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
is False
)
@pytest.mark.asyncio
async def test_concurrent_grant_revoke(
self,
permission_manager: PermissionManager,
) -> None:
"""Test concurrent grant and revoke operations."""
async def grant_many():
grants = []
for i in range(10):
g = await permission_manager.grant(
agent_id="agent-1",
resource_pattern=f"/path{i}/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
grants.append(g)
return grants
async def revoke_many(grants):
for g in grants:
await permission_manager.revoke(g.id)
grants = await grant_many()
await revoke_many(grants)
# All should be revoked
remaining = await permission_manager.list_grants(agent_id="agent-1")
assert len(remaining) == 0
@pytest.mark.asyncio
async def test_check_action_with_no_resource_or_tool(
self,
permission_manager: PermissionManager,
action_metadata: ActionMetadata,
) -> None:
"""Test check_action when both resource and tool_name are None."""
await permission_manager.grant(
agent_id="test-agent",
resource_pattern="*",
resource_type=ResourceType.LLM,
level=PermissionLevel.EXECUTE,
)
action = ActionRequest(
action_type=ActionType.LLM_CALL,
resource=None,
tool_name=None,
metadata=action_metadata,
)
# Should use "*" as fallback
allowed = await permission_manager.check_action(action)
assert allowed is True
@pytest.mark.asyncio
async def test_cleanup_expired_called_on_check(
self,
permission_manager: PermissionManager,
) -> None:
"""Test that expired grants are cleaned up during check."""
# Create expired grant
grant = await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/old/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
duration_seconds=1,
)
grant.expires_at = datetime.utcnow() - timedelta(seconds=10)
# Create valid grant
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/new/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Run a check - this should trigger cleanup
await permission_manager.check(
agent_id="agent-1",
resource="/new/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
# Now verify expired grant was cleaned up
async with permission_manager._lock:
assert len(permission_manager._grants) == 1
assert permission_manager._grants[0].resource_pattern == "/new/*"
@pytest.mark.asyncio
async def test_check_wrong_agent_id(
self,
permission_manager: PermissionManager,
) -> None:
"""Test check fails for different agent."""
await permission_manager.grant(
agent_id="agent-1",
resource_pattern="/data/*",
resource_type=ResourceType.FILE,
level=PermissionLevel.READ,
)
# Different agent should not have access
allowed = await permission_manager.check(
agent_id="agent-2",
resource="/data/file.txt",
resource_type=ResourceType.FILE,
required_level=PermissionLevel.READ,
)
assert allowed is False

823
backend/tests/services/safety/test_rollback.py Normal file
View File

@@ -0,0 +1,823 @@
"""Tests for Rollback Manager.
Tests cover:
- FileCheckpoint: state storage
- RollbackManager: checkpoint, rollback, cleanup
- TransactionContext: auto-rollback, commit, manual rollback
- Edge cases: non-existent files, partial failures, expiration
"""
import tempfile
from datetime import datetime, timedelta
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
import pytest_asyncio
from app.services.safety.exceptions import RollbackError
from app.services.safety.models import (
ActionMetadata,
ActionRequest,
ActionType,
CheckpointType,
)
from app.services.safety.rollback.manager import (
FileCheckpoint,
RollbackManager,
TransactionContext,
)
# ============================================================================
# Fixtures
# ============================================================================
@pytest.fixture
def action_metadata() -> ActionMetadata:
"""Create standard action metadata for tests."""
return ActionMetadata(
agent_id="test-agent",
project_id="test-project",
session_id="test-session",
)
@pytest.fixture
def action_request(action_metadata: ActionMetadata) -> ActionRequest:
"""Create a standard action request for tests."""
return ActionRequest(
id="action-123",
action_type=ActionType.FILE_WRITE,
tool_name="file_write",
resource="/tmp/test_file.txt", # noqa: S108
metadata=action_metadata,
is_destructive=True,
)
@pytest_asyncio.fixture
async def rollback_manager() -> RollbackManager:
"""Create a RollbackManager for testing."""
with tempfile.TemporaryDirectory() as tmpdir:
with patch("app.services.safety.rollback.manager.get_safety_config") as mock:
mock.return_value = MagicMock(
checkpoint_dir=tmpdir,
checkpoint_retention_hours=24,
)
manager = RollbackManager(checkpoint_dir=tmpdir, retention_hours=24)
yield manager
@pytest.fixture
def temp_dir() -> Path:
"""Create a temporary directory for file operations."""
with tempfile.TemporaryDirectory() as tmpdir:
yield Path(tmpdir)
# ============================================================================
# FileCheckpoint Tests
# ============================================================================
class TestFileCheckpoint:
"""Tests for the FileCheckpoint class."""
def test_file_checkpoint_creation(self) -> None:
"""Test creating a file checkpoint."""
fc = FileCheckpoint(
checkpoint_id="cp-123",
file_path="/path/to/file.txt",
original_content=b"original content",
existed=True,
)
assert fc.checkpoint_id == "cp-123"
assert fc.file_path == "/path/to/file.txt"
assert fc.original_content == b"original content"
assert fc.existed is True
assert fc.created_at is not None
def test_file_checkpoint_nonexistent_file(self) -> None:
"""Test checkpoint for non-existent file."""
fc = FileCheckpoint(
checkpoint_id="cp-123",
file_path="/path/to/new_file.txt",
original_content=None,
existed=False,
)
assert fc.original_content is None
assert fc.existed is False
# ============================================================================
# RollbackManager Tests
# ============================================================================
class TestRollbackManager:
"""Tests for the RollbackManager class."""
@pytest.mark.asyncio
async def test_create_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test creating a checkpoint."""
checkpoint = await rollback_manager.create_checkpoint(
action=action_request,
checkpoint_type=CheckpointType.FILE,
description="Test checkpoint",
)
assert checkpoint.id is not None
assert checkpoint.action_id == action_request.id
assert checkpoint.checkpoint_type == CheckpointType.FILE
assert checkpoint.description == "Test checkpoint"
assert checkpoint.expires_at is not None
assert checkpoint.is_valid is True
@pytest.mark.asyncio
async def test_create_checkpoint_default_description(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test checkpoint with default description."""
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
assert "file_write" in checkpoint.description
@pytest.mark.asyncio
async def test_checkpoint_file_exists(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpointing an existing file."""
# Create a file
test_file = temp_dir / "test.txt"
test_file.write_text("original content")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Verify checkpoint was stored
async with rollback_manager._lock:
file_checkpoints = rollback_manager._file_checkpoints.get(checkpoint.id, [])
assert len(file_checkpoints) == 1
assert file_checkpoints[0].existed is True
assert file_checkpoints[0].original_content == b"original content"
@pytest.mark.asyncio
async def test_checkpoint_file_not_exists(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpointing a non-existent file."""
test_file = temp_dir / "new_file.txt"
assert not test_file.exists()
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Verify checkpoint was stored
async with rollback_manager._lock:
file_checkpoints = rollback_manager._file_checkpoints.get(checkpoint.id, [])
assert len(file_checkpoints) == 1
assert file_checkpoints[0].existed is False
assert file_checkpoints[0].original_content is None
@pytest.mark.asyncio
async def test_checkpoint_files_multiple(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpointing multiple files."""
# Create files
file1 = temp_dir / "file1.txt"
file2 = temp_dir / "file2.txt"
file1.write_text("content 1")
file2.write_text("content 2")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_files(
checkpoint.id,
[str(file1), str(file2)],
)
async with rollback_manager._lock:
file_checkpoints = rollback_manager._file_checkpoints.get(checkpoint.id, [])
assert len(file_checkpoints) == 2
@pytest.mark.asyncio
async def test_rollback_restore_modified_file(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test rollback restores modified file content."""
test_file = temp_dir / "test.txt"
test_file.write_text("original content")
# Create checkpoint
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Modify file
test_file.write_text("modified content")
assert test_file.read_text() == "modified content"
# Rollback
result = await rollback_manager.rollback(checkpoint.id)
assert result.success is True
assert len(result.actions_rolled_back) == 1
assert test_file.read_text() == "original content"
@pytest.mark.asyncio
async def test_rollback_delete_new_file(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test rollback deletes file that didn't exist before."""
test_file = temp_dir / "new_file.txt"
assert not test_file.exists()
# Create checkpoint before file exists
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Create the file
test_file.write_text("new content")
assert test_file.exists()
# Rollback
result = await rollback_manager.rollback(checkpoint.id)
assert result.success is True
assert not test_file.exists()
@pytest.mark.asyncio
async def test_rollback_not_found(
self,
rollback_manager: RollbackManager,
) -> None:
"""Test rollback with non-existent checkpoint."""
with pytest.raises(RollbackError) as exc_info:
await rollback_manager.rollback("nonexistent-id")
assert "not found" in str(exc_info.value)
@pytest.mark.asyncio
async def test_rollback_invalid_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test rollback with invalidated checkpoint."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Rollback once (invalidates checkpoint)
await rollback_manager.rollback(checkpoint.id)
# Try to rollback again
with pytest.raises(RollbackError) as exc_info:
await rollback_manager.rollback(checkpoint.id)
assert "no longer valid" in str(exc_info.value)
@pytest.mark.asyncio
async def test_discard_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test discarding a checkpoint."""
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
result = await rollback_manager.discard_checkpoint(checkpoint.id)
assert result is True
# Verify it's gone
cp = await rollback_manager.get_checkpoint(checkpoint.id)
assert cp is None
@pytest.mark.asyncio
async def test_discard_checkpoint_nonexistent(
self,
rollback_manager: RollbackManager,
) -> None:
"""Test discarding a non-existent checkpoint."""
result = await rollback_manager.discard_checkpoint("nonexistent-id")
assert result is False
@pytest.mark.asyncio
async def test_get_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test getting a checkpoint by ID."""
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
retrieved = await rollback_manager.get_checkpoint(checkpoint.id)
assert retrieved is not None
assert retrieved.id == checkpoint.id
@pytest.mark.asyncio
async def test_get_checkpoint_nonexistent(
self,
rollback_manager: RollbackManager,
) -> None:
"""Test getting a non-existent checkpoint."""
retrieved = await rollback_manager.get_checkpoint("nonexistent-id")
assert retrieved is None
@pytest.mark.asyncio
async def test_list_checkpoints(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test listing checkpoints."""
await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.create_checkpoint(action=action_request)
checkpoints = await rollback_manager.list_checkpoints()
assert len(checkpoints) == 2
@pytest.mark.asyncio
async def test_list_checkpoints_by_action(
self,
rollback_manager: RollbackManager,
action_metadata: ActionMetadata,
) -> None:
"""Test listing checkpoints filtered by action."""
action1 = ActionRequest(
id="action-1",
action_type=ActionType.FILE_WRITE,
metadata=action_metadata,
)
action2 = ActionRequest(
id="action-2",
action_type=ActionType.FILE_WRITE,
metadata=action_metadata,
)
await rollback_manager.create_checkpoint(action=action1)
await rollback_manager.create_checkpoint(action=action2)
checkpoints = await rollback_manager.list_checkpoints(action_id="action-1")
assert len(checkpoints) == 1
assert checkpoints[0].action_id == "action-1"
@pytest.mark.asyncio
async def test_list_checkpoints_excludes_expired(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test list_checkpoints excludes expired by default."""
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
# Manually expire it
async with rollback_manager._lock:
rollback_manager._checkpoints[checkpoint.id].expires_at = (
datetime.utcnow() - timedelta(hours=1)
)
checkpoints = await rollback_manager.list_checkpoints()
assert len(checkpoints) == 0
# With include_expired=True
checkpoints = await rollback_manager.list_checkpoints(include_expired=True)
assert len(checkpoints) == 1
@pytest.mark.asyncio
async def test_cleanup_expired(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test cleaning up expired checkpoints."""
# Create checkpoints
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
test_file = temp_dir / "test.txt"
test_file.write_text("content")
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Expire it
async with rollback_manager._lock:
rollback_manager._checkpoints[checkpoint.id].expires_at = (
datetime.utcnow() - timedelta(hours=1)
)
# Cleanup
count = await rollback_manager.cleanup_expired()
assert count == 1
# Verify it's gone
async with rollback_manager._lock:
assert checkpoint.id not in rollback_manager._checkpoints
assert checkpoint.id not in rollback_manager._file_checkpoints
# ============================================================================
# TransactionContext Tests
# ============================================================================
class TestTransactionContext:
"""Tests for the TransactionContext class."""
@pytest.mark.asyncio
async def test_context_creates_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test that entering context creates a checkpoint."""
async with TransactionContext(rollback_manager, action_request) as tx:
assert tx.checkpoint_id is not None
# Verify checkpoint exists
cp = await rollback_manager.get_checkpoint(tx.checkpoint_id)
assert cp is not None
@pytest.mark.asyncio
async def test_context_checkpoint_file(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpointing files through context."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
async with TransactionContext(rollback_manager, action_request) as tx:
await tx.checkpoint_file(str(test_file))
# Modify file
test_file.write_text("modified")
# Manual rollback
result = await tx.rollback()
assert result is not None
assert result.success is True
assert test_file.read_text() == "original"
@pytest.mark.asyncio
async def test_context_checkpoint_files(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpointing multiple files through context."""
file1 = temp_dir / "file1.txt"
file2 = temp_dir / "file2.txt"
file1.write_text("content 1")
file2.write_text("content 2")
async with TransactionContext(rollback_manager, action_request) as tx:
await tx.checkpoint_files([str(file1), str(file2)])
cp_id = tx.checkpoint_id
async with rollback_manager._lock:
file_cps = rollback_manager._file_checkpoints.get(cp_id, [])
assert len(file_cps) == 2
tx.commit()
@pytest.mark.asyncio
async def test_context_auto_rollback_on_exception(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test auto-rollback when exception occurs."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
with pytest.raises(ValueError):
async with TransactionContext(rollback_manager, action_request) as tx:
await tx.checkpoint_file(str(test_file))
test_file.write_text("modified")
raise ValueError("Simulated error")
# Should have been rolled back
assert test_file.read_text() == "original"
@pytest.mark.asyncio
async def test_context_commit_prevents_rollback(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test that commit prevents auto-rollback."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
with pytest.raises(ValueError):
async with TransactionContext(rollback_manager, action_request) as tx:
await tx.checkpoint_file(str(test_file))
test_file.write_text("modified")
tx.commit()
raise ValueError("Simulated error after commit")
# Should NOT have been rolled back
assert test_file.read_text() == "modified"
@pytest.mark.asyncio
async def test_context_discards_checkpoint_on_commit(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test that checkpoint is discarded after successful commit."""
checkpoint_id = None
async with TransactionContext(rollback_manager, action_request) as tx:
checkpoint_id = tx.checkpoint_id
tx.commit()
# Checkpoint should be discarded
cp = await rollback_manager.get_checkpoint(checkpoint_id)
assert cp is None
@pytest.mark.asyncio
async def test_context_no_auto_rollback_when_disabled(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test that auto_rollback=False disables auto-rollback."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
with pytest.raises(ValueError):
async with TransactionContext(
rollback_manager,
action_request,
auto_rollback=False,
) as tx:
await tx.checkpoint_file(str(test_file))
test_file.write_text("modified")
raise ValueError("Simulated error")
# Should NOT have been rolled back
assert test_file.read_text() == "modified"
@pytest.mark.asyncio
async def test_context_manual_rollback(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test manual rollback within context."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
async with TransactionContext(rollback_manager, action_request) as tx:
await tx.checkpoint_file(str(test_file))
test_file.write_text("modified")
# Manual rollback
result = await tx.rollback()
assert result is not None
assert result.success is True
assert test_file.read_text() == "original"
@pytest.mark.asyncio
async def test_context_rollback_without_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test rollback when checkpoint is None."""
tx = TransactionContext(rollback_manager, action_request)
# Don't enter context, so _checkpoint is None
result = await tx.rollback()
assert result is None
@pytest.mark.asyncio
async def test_context_checkpoint_file_without_checkpoint(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test checkpoint_file when checkpoint is None (no-op)."""
tx = TransactionContext(rollback_manager, action_request)
test_file = temp_dir / "test.txt"
test_file.write_text("content")
# Should not raise - just a no-op
await tx.checkpoint_file(str(test_file))
await tx.checkpoint_files([str(test_file)])
# ============================================================================
# Edge Cases
# ============================================================================
class TestRollbackEdgeCases:
"""Edge cases that could reveal hidden bugs."""
@pytest.mark.asyncio
async def test_checkpoint_file_for_unknown_checkpoint(
self,
rollback_manager: RollbackManager,
temp_dir: Path,
) -> None:
"""Test checkpointing file for non-existent checkpoint."""
test_file = temp_dir / "test.txt"
test_file.write_text("content")
# Should create the list if it doesn't exist
await rollback_manager.checkpoint_file("unknown-checkpoint", str(test_file))
async with rollback_manager._lock:
assert "unknown-checkpoint" in rollback_manager._file_checkpoints
@pytest.mark.asyncio
async def test_rollback_with_partial_failure(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test rollback when some files fail to restore."""
file1 = temp_dir / "file1.txt"
file1.write_text("original 1")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(file1))
# Add a file checkpoint with a path that will fail
async with rollback_manager._lock:
# Create a checkpoint for a file in a non-writable location
bad_fc = FileCheckpoint(
checkpoint_id=checkpoint.id,
file_path="/nonexistent/path/file.txt",
original_content=b"content",
existed=True,
)
rollback_manager._file_checkpoints[checkpoint.id].append(bad_fc)
# Rollback - partial failure expected
result = await rollback_manager.rollback(checkpoint.id)
assert result.success is False
assert len(result.actions_rolled_back) == 1
assert len(result.failed_actions) == 1
assert "Failed to rollback" in result.error
@pytest.mark.asyncio
async def test_rollback_file_creates_parent_dirs(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test that rollback creates parent directories if needed."""
nested_file = temp_dir / "subdir" / "nested" / "file.txt"
nested_file.parent.mkdir(parents=True)
nested_file.write_text("original")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(nested_file))
# Delete the entire directory structure
nested_file.unlink()
(temp_dir / "subdir" / "nested").rmdir()
(temp_dir / "subdir").rmdir()
# Rollback should recreate
result = await rollback_manager.rollback(checkpoint.id)
assert result.success is True
assert nested_file.exists()
assert nested_file.read_text() == "original"
@pytest.mark.asyncio
async def test_rollback_file_already_correct(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test rollback when file already has correct content."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
await rollback_manager.checkpoint_file(checkpoint.id, str(test_file))
# Don't modify file - rollback should still succeed
result = await rollback_manager.rollback(checkpoint.id)
assert result.success is True
assert test_file.read_text() == "original"
@pytest.mark.asyncio
async def test_checkpoint_with_none_expires_at(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test list_checkpoints handles None expires_at."""
checkpoint = await rollback_manager.create_checkpoint(action=action_request)
# Set expires_at to None
async with rollback_manager._lock:
rollback_manager._checkpoints[checkpoint.id].expires_at = None
# Should still be listed
checkpoints = await rollback_manager.list_checkpoints()
assert len(checkpoints) == 1
@pytest.mark.asyncio
async def test_auto_rollback_failure_logged(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
temp_dir: Path,
) -> None:
"""Test that auto-rollback failure is logged, not raised."""
test_file = temp_dir / "test.txt"
test_file.write_text("original")
with patch.object(
rollback_manager, "rollback", side_effect=Exception("Rollback failed!")
):
with patch("app.services.safety.rollback.manager.logger") as mock_logger:
with pytest.raises(ValueError):
async with TransactionContext(
rollback_manager, action_request
) as tx:
await tx.checkpoint_file(str(test_file))
test_file.write_text("modified")
raise ValueError("Original error")
# Rollback error should be logged
mock_logger.error.assert_called()
@pytest.mark.asyncio
async def test_multiple_checkpoints_same_action(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test creating multiple checkpoints for the same action."""
cp1 = await rollback_manager.create_checkpoint(action=action_request)
cp2 = await rollback_manager.create_checkpoint(action=action_request)
assert cp1.id != cp2.id
checkpoints = await rollback_manager.list_checkpoints(
action_id=action_request.id
)
assert len(checkpoints) == 2
@pytest.mark.asyncio
async def test_cleanup_expired_with_no_expired(
self,
rollback_manager: RollbackManager,
action_request: ActionRequest,
) -> None:
"""Test cleanup when no checkpoints are expired."""
await rollback_manager.create_checkpoint(action=action_request)
count = await rollback_manager.cleanup_expired()
assert count == 0
# Checkpoint should still exist
checkpoints = await rollback_manager.list_checkpoints()
assert len(checkpoints) == 1

359
backend/tests/services/safety/test_validation.py
View File

@@ -363,6 +363,365 @@ class TestValidationBatch:
assert results[1].decision == SafetyDecision.DENY
class TestValidationCache:
"""Tests for ValidationCache class."""
@pytest.mark.asyncio
async def test_cache_get_miss(self) -> None:
"""Test cache miss."""
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=10, ttl_seconds=60)
result = await cache.get("nonexistent")
assert result is None
@pytest.mark.asyncio
async def test_cache_get_hit(self) -> None:
"""Test cache hit."""
from app.services.safety.models import ValidationResult
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=10, ttl_seconds=60)
vr = ValidationResult(
action_id="action-1",
decision=SafetyDecision.ALLOW,
applied_rules=[],
reasons=["test"],
)
await cache.set("key1", vr)
result = await cache.get("key1")
assert result is not None
assert result.action_id == "action-1"
@pytest.mark.asyncio
async def test_cache_ttl_expiry(self) -> None:
"""Test cache TTL expiry."""
import time
from unittest.mock import patch
from app.services.safety.models import ValidationResult
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=10, ttl_seconds=1)
vr = ValidationResult(
action_id="action-1",
decision=SafetyDecision.ALLOW,
applied_rules=[],
reasons=["test"],
)
await cache.set("key1", vr)
# Advance time past TTL
with patch("time.time", return_value=time.time() + 2):
result = await cache.get("key1")
assert result is None # Should be expired
@pytest.mark.asyncio
async def test_cache_eviction_on_full(self) -> None:
"""Test cache eviction when full."""
from app.services.safety.models import ValidationResult
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=2, ttl_seconds=60)
vr1 = ValidationResult(action_id="a1", decision=SafetyDecision.ALLOW)
vr2 = ValidationResult(action_id="a2", decision=SafetyDecision.ALLOW)
vr3 = ValidationResult(action_id="a3", decision=SafetyDecision.ALLOW)
await cache.set("key1", vr1)
await cache.set("key2", vr2)
await cache.set("key3", vr3) # Should evict key1
# key1 should be evicted
assert await cache.get("key1") is None
assert await cache.get("key2") is not None
assert await cache.get("key3") is not None
@pytest.mark.asyncio
async def test_cache_update_existing_key(self) -> None:
"""Test updating existing key in cache."""
from app.services.safety.models import ValidationResult
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=10, ttl_seconds=60)
vr1 = ValidationResult(action_id="a1", decision=SafetyDecision.ALLOW)
vr2 = ValidationResult(action_id="a1-updated", decision=SafetyDecision.DENY)
await cache.set("key1", vr1)
await cache.set("key1", vr2) # Should update, not add
result = await cache.get("key1")
assert result is not None
assert result.action_id == "a1" # Still old value since we move_to_end
@pytest.mark.asyncio
async def test_cache_clear(self) -> None:
"""Test clearing cache."""
from app.services.safety.models import ValidationResult
from app.services.safety.validation.validator import ValidationCache
cache = ValidationCache(max_size=10, ttl_seconds=60)
vr = ValidationResult(action_id="a1", decision=SafetyDecision.ALLOW)
await cache.set("key1", vr)
await cache.set("key2", vr)
await cache.clear()
assert await cache.get("key1") is None
assert await cache.get("key2") is None
class TestValidatorCaching:
"""Tests for validator caching functionality."""
@pytest.mark.asyncio
async def test_cache_hit(self) -> None:
"""Test that cache is used for repeated validations."""
validator = ActionValidator(cache_enabled=True, cache_ttl=60)
metadata = ActionMetadata(agent_id="test-agent", session_id="session-1")
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name="file_read",
resource="/tmp/test.txt", # noqa: S108
metadata=metadata,
)
# First call populates cache
result1 = await validator.validate(action)
# Second call should use cache
result2 = await validator.validate(action)
assert result1.decision == result2.decision
@pytest.mark.asyncio
async def test_clear_cache(self) -> None:
"""Test clearing the validation cache."""
validator = ActionValidator(cache_enabled=True)
metadata = ActionMetadata(agent_id="test-agent", session_id="session-1")
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name="file_read",
metadata=metadata,
)
await validator.validate(action)
await validator.clear_cache()
# Cache should be empty now (no error)
result = await validator.validate(action)
assert result.decision == SafetyDecision.ALLOW
class TestRuleMatching:
"""Tests for rule matching edge cases."""
@pytest.mark.asyncio
async def test_action_type_mismatch(self) -> None:
"""Test that rule doesn't match when action type doesn't match."""
validator = ActionValidator(cache_enabled=False)
validator.add_rule(
ValidationRule(
name="file_only",
action_types=[ActionType.FILE_READ],
decision=SafetyDecision.DENY,
)
)
metadata = ActionMetadata(agent_id="test-agent")
action = ActionRequest(
action_type=ActionType.SHELL_COMMAND, # Different type
tool_name="shell_exec",
metadata=metadata,
)
result = await validator.validate(action)
assert result.decision == SafetyDecision.ALLOW # Rule didn't match
@pytest.mark.asyncio
async def test_tool_pattern_no_tool_name(self) -> None:
"""Test rule with tool pattern when action has no tool_name."""
validator = ActionValidator(cache_enabled=False)
validator.add_rule(
create_deny_rule(
name="deny_files",
tool_patterns=["file_*"],
)
)
metadata = ActionMetadata(agent_id="test-agent")
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name=None, # No tool name
metadata=metadata,
)
result = await validator.validate(action)
assert result.decision == SafetyDecision.ALLOW # Rule didn't match
@pytest.mark.asyncio
async def test_resource_pattern_no_resource(self) -> None:
"""Test rule with resource pattern when action has no resource."""
validator = ActionValidator(cache_enabled=False)
validator.add_rule(
create_deny_rule(
name="deny_secrets",
resource_patterns=["/secret/*"],
)
)
metadata = ActionMetadata(agent_id="test-agent")
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name="file_read",
resource=None, # No resource
metadata=metadata,
)
result = await validator.validate(action)
assert result.decision == SafetyDecision.ALLOW # Rule didn't match
@pytest.mark.asyncio
async def test_resource_pattern_no_match(self) -> None:
"""Test rule with resource pattern that doesn't match."""
validator = ActionValidator(cache_enabled=False)
validator.add_rule(
create_deny_rule(
name="deny_secrets",
resource_patterns=["/secret/*"],
)
)
metadata = ActionMetadata(agent_id="test-agent")
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name="file_read",
resource="/public/file.txt", # Doesn't match
metadata=metadata,
)
result = await validator.validate(action)
assert result.decision == SafetyDecision.ALLOW # Pattern didn't match
class TestPolicyLoading:
"""Tests for policy loading edge cases."""
@pytest.mark.asyncio
async def test_load_rules_from_policy_with_validation_rules(self) -> None:
"""Test loading policy with explicit validation rules."""
validator = ActionValidator(cache_enabled=False)
rule = ValidationRule(
name="policy_rule",
tool_patterns=["test_*"],
decision=SafetyDecision.DENY,
reason="From policy",
)
policy = SafetyPolicy(
name="test",
validation_rules=[rule],
require_approval_for=[], # Clear defaults
denied_tools=[], # Clear defaults
)
validator.load_rules_from_policy(policy)
assert len(validator._rules) == 1
assert validator._rules[0].name == "policy_rule"
@pytest.mark.asyncio
async def test_load_approval_all_pattern(self) -> None:
"""Test loading policy with * approval pattern (all actions)."""
validator = ActionValidator(cache_enabled=False)
policy = SafetyPolicy(
name="test",
require_approval_for=["*"], # All actions require approval
denied_tools=[], # Clear defaults
)
validator.load_rules_from_policy(policy)
approval_rules = [
r for r in validator._rules if r.decision == SafetyDecision.REQUIRE_APPROVAL
]
assert len(approval_rules) == 1
assert approval_rules[0].name == "require_approval_all"
assert approval_rules[0].action_types == list(ActionType)
@pytest.mark.asyncio
async def test_validate_with_policy_loads_rules(self) -> None:
"""Test that validate() loads rules from policy if none exist."""
validator = ActionValidator(cache_enabled=False)
policy = SafetyPolicy(
name="test",
denied_tools=["dangerous_*"],
)
metadata = ActionMetadata(agent_id="test-agent")
action = ActionRequest(
action_type=ActionType.SHELL_COMMAND,
tool_name="dangerous_exec",
metadata=metadata,
)
# Validate with policy - should load rules
result = await validator.validate(action, policy=policy)
assert result.decision == SafetyDecision.DENY
class TestCacheKeyGeneration:
"""Tests for cache key generation."""
def test_get_cache_key(self) -> None:
"""Test cache key generation."""
validator = ActionValidator(cache_enabled=True)
metadata = ActionMetadata(
agent_id="test-agent",
autonomy_level=AutonomyLevel.MILESTONE,
)
action = ActionRequest(
action_type=ActionType.FILE_READ,
tool_name="file_read",
resource="/tmp/test.txt", # noqa: S108
metadata=metadata,
)
key = validator._get_cache_key(action)
assert "file_read" in key
assert "file_read" in key
assert "/tmp/test.txt" in key # noqa: S108
assert "test-agent" in key
assert "milestone" in key
def test_get_cache_key_no_resource(self) -> None:
"""Test cache key generation without resource."""
validator = ActionValidator(cache_enabled=True)
metadata = ActionMetadata(agent_id="agent-1")
action = ActionRequest(
action_type=ActionType.SHELL_COMMAND,
tool_name="shell_exec",
resource=None,
metadata=metadata,
)
key = validator._get_cache_key(action)
# Should not error with None resource
assert "shell" in key
assert "agent-1" in key
class TestHelperFunctions:
"""Tests for rule creation helper functions."""