arbiter/tests/test_models.py

"""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"