"""Tests for data models.""" from pathlib import Path import pytest from arbiter.models import ( AgentConfig, AgentName, Finding, Policy, ReviewResult, Severity, Verdict, ) class TestEnums: def test_severity_values(self) -> None: assert Severity.CRITICAL == "critical" assert Severity.HIGH == "high" assert Severity.MEDIUM == "medium" assert Severity.LOW == "low" assert Severity.INFO == "info" def test_verdict_values(self) -> None: assert Verdict.APPROVE == "approve" assert Verdict.REQUEST_CHANGES == "request_changes" assert Verdict.COMMENT == "comment" def test_agent_name_values(self) -> None: assert AgentName.SECURITY == "security" assert AgentName.STYLE == "style" assert AgentName.COMPLEXITY == "complexity" def test_severity_from_string(self) -> None: assert Severity("critical") == Severity.CRITICAL assert Severity("high") == Severity.HIGH class TestFinding: def test_finding_creation(self) -> None: finding = Finding( id="test-123", agent=AgentName.SECURITY, file="src/auth.py", line_start=10, line_end=15, severity=Severity.HIGH, confidence=0.9, title="SQL Injection", description="User input concatenated in SQL query", reasoning="Allows attackers to execute arbitrary SQL", suggestion="Use parameterized queries", references=["https://owasp.org"], prompt_version="security-v1.0", ) assert finding.id == "test-123" assert finding.agent == AgentName.SECURITY assert finding.severity == Severity.HIGH assert finding.confidence == 0.9 def test_finding_confidence_validation(self) -> None: with pytest.raises(ValueError): Finding( id="test", agent=AgentName.SECURITY, file="test.py", line_start=1, line_end=1, severity=Severity.INFO, confidence=1.5, title="Test", description="Test", reasoning="Test", prompt_version="test-v1.0", ) def test_finding_line_validation(self) -> None: with pytest.raises(ValueError): Finding( id="test", agent=AgentName.SECURITY, file="test.py", line_start=0, line_end=1, severity=Severity.INFO, confidence=0.5, title="Test", description="Test", reasoning="Test", prompt_version="test-v1.0", ) def test_finding_serialization(self) -> None: finding = Finding( id="test-123", agent=AgentName.SECURITY, file="src/auth.py", line_start=10, line_end=15, severity=Severity.HIGH, confidence=0.9, title="Test", description="Test desc", reasoning="Test reason", prompt_version="security-v1.0", ) data = finding.model_dump() assert data["id"] == "test-123" assert data["agent"] == "security" assert data["severity"] == "high" class TestAgentConfig: def test_default_values(self) -> None: config = AgentConfig() assert config.enabled is True assert config.model is None assert config.severity_threshold == Severity.INFO assert config.prompt_additions is None def test_custom_values(self) -> None: config = AgentConfig( enabled=False, model="gpt-4o", severity_threshold=Severity.MEDIUM, prompt_additions="Focus on auth", ) assert config.enabled is False assert config.model == "gpt-4o" assert config.severity_threshold == Severity.MEDIUM class TestPolicy: def test_default_policy(self) -> None: policy = Policy() assert len(policy.agents) == 3 assert AgentName.SECURITY in policy.agents assert AgentName.STYLE in policy.agents assert AgentName.COMPLEXITY in policy.agents def test_get_enabled_agents(self) -> None: policy = Policy() enabled = policy.get_enabled_agents() assert len(enabled) == 3 assert AgentName.SECURITY in enabled def test_get_enabled_agents_with_disabled(self) -> None: policy = Policy( agents={ AgentName.SECURITY: AgentConfig(enabled=True), AgentName.STYLE: AgentConfig(enabled=False), AgentName.COMPLEXITY: AgentConfig(enabled=True), } ) enabled = policy.get_enabled_agents() assert len(enabled) == 2 assert AgentName.STYLE not in enabled def test_load_from_yaml(self, tmp_path: Path) -> None: policy_file = tmp_path / "policy.yaml" policy_file.write_text(""" version: "1.1" agents: security: enabled: true model: gpt-4o severity_threshold: high style: enabled: false complexity: enabled: true """) policy = Policy.load(policy_file) assert policy.version == "1.1" assert policy.agents[AgentName.SECURITY].model == "gpt-4o" assert policy.agents[AgentName.SECURITY].severity_threshold == Severity.HIGH assert policy.agents[AgentName.STYLE].enabled is False def test_load_empty_yaml(self, tmp_path: Path) -> None: policy_file = tmp_path / "policy.yaml" policy_file.write_text("") policy = Policy.load(policy_file) assert policy.version == "1.0" assert len(policy.agents) == 3 class TestReviewResult: def test_review_result_creation(self) -> None: result = ReviewResult( agent_name=AgentName.SECURITY, findings=[], duration_ms=1000, tokens_used=500, cost_usd=0.01, ) assert result.agent_name == AgentName.SECURITY assert result.duration_ms == 1000 assert result.cost_usd == 0.01 def test_review_result_with_findings(self) -> None: finding = Finding( id="test-123", agent=AgentName.SECURITY, file="test.py", line_start=1, line_end=1, severity=Severity.HIGH, confidence=0.9, title="Test", description="Test", reasoning="Test", prompt_version="test-v1.0", ) result = ReviewResult( agent_name=AgentName.SECURITY, findings=[finding], duration_ms=1000, tokens_used=500, cost_usd=0.01, ) assert len(result.findings) == 1 assert result.findings[0].id == "test-123"