tests/test_h1_draft.py

"""Tests for sentinel/services/h1_draft.py — H1 bug report draft generator.

The LLM backend is mocked throughout so no network calls are made.
"""
from __future__ import annotations

from unittest.mock import AsyncMock, patch

import pytest

from sentinel.services.h1_draft import generate_h1_draft, _H1_PROMPT


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _finding(overrides: dict | None = None) -> dict:
    base = {
        "tool": "agent-audit",
        "rule": "AGENT-010",
        "severity": "ERROR",
        "confidence": "confirmed",
        "file": "app.py",
        "line": 42,
        "message": "User-controlled input concatenated into system prompt",
        "owasp": ["A03:2021-Injection"],
        "remediation": "Sanitize or reject user input before including in system prompt",
        "score": 8,
    }
    if overrides:
        base.update(overrides)
    return base


_MOCK_DRAFT = "### Title\nPrompt Injection in app.py\n\n### Summary\nTest draft."


# ---------------------------------------------------------------------------
# Prompt template
# ---------------------------------------------------------------------------
class TestH1Prompt:
    def test_template_has_required_placeholders(self):
        for key in ("{target}", "{tool}", "{rule}", "{severity}", "{confidence}",
                    "{file}", "{line}", "{message}", "{owasp}", "{remediation}"):
            assert key in _H1_PROMPT, f"Missing placeholder: {key}"

    def test_template_renders_without_error(self):
        rendered = _H1_PROMPT.format(
            target="https://huggingface.co/spaces/test/space",
            tool="bandit", rule="B301", severity="ERROR", confidence="confirmed",
            file="app.py", line=10, message="pickle.load() called", owasp="A08",
            remediation="Use json instead of pickle",
        )
        assert "bandit" in rendered
        assert "B301" in rendered
        assert "pickle.load()" in rendered


# ---------------------------------------------------------------------------
# generate_h1_draft — happy path
# ---------------------------------------------------------------------------
class TestGenerateH1Draft:
    async def test_returns_llm_response(self):
        with patch("sentinel.services.ai_explain.call_llm", new=AsyncMock(return_value=_MOCK_DRAFT)):
            result = await generate_h1_draft(_finding(), scan=None, target="https://huggingface.co/spaces/test/space")
        assert result == _MOCK_DRAFT

    async def test_prompt_includes_finding_fields(self):
        captured = {}

        async def _mock_llm(prompt, max_tokens=512):
            captured["prompt"] = prompt
            return _MOCK_DRAFT

        with patch("sentinel.services.ai_explain.call_llm", new=_mock_llm):
            await generate_h1_draft(_finding(), scan=None, target="my-target-space")

        prompt = captured["prompt"]
        assert "agent-audit" in prompt
        assert "AGENT-010" in prompt
        assert "app.py" in prompt
        assert "my-target-space" in prompt
        assert "42" in prompt

    async def test_owasp_list_joined(self):
        captured = {}

        async def _mock_llm(prompt, max_tokens=512):
            captured["prompt"] = prompt
            return _MOCK_DRAFT

        f = _finding({"owasp": ["A03:2021-Injection", "A06:2021-Outdated"]})
        with patch("sentinel.services.ai_explain.call_llm", new=_mock_llm):
            await generate_h1_draft(f, scan=None, target="test")

        assert "A03:2021-Injection" in captured["prompt"]
        assert "A06:2021-Outdated" in captured["prompt"]

    async def test_owasp_string_passed_through(self):
        captured = {}

        async def _mock_llm(prompt, max_tokens=512):
            captured["prompt"] = prompt
            return _MOCK_DRAFT

        f = _finding({"owasp": "A03:2021-Injection"})
        with patch("sentinel.services.ai_explain.call_llm", new=_mock_llm):
            await generate_h1_draft(f, scan=None, target="test")

        assert "A03:2021-Injection" in captured["prompt"]

    async def test_max_tokens_800(self):
        captured = {}

        async def _mock_llm(prompt, max_tokens=512):
            captured["max_tokens"] = max_tokens
            return _MOCK_DRAFT

        with patch("sentinel.services.ai_explain.call_llm", new=_mock_llm):
            await generate_h1_draft(_finding(), scan=None, target="test")

        assert captured["max_tokens"] == 800

    async def test_message_truncated_to_400_chars(self):
        captured = {}

        async def _mock_llm(prompt, max_tokens=512):
            captured["prompt"] = prompt
            return _MOCK_DRAFT

        long_message = "X" * 600
        f = _finding({"message": long_message})
        with patch("sentinel.services.ai_explain.call_llm", new=_mock_llm):
            await generate_h1_draft(f, scan=None, target="test")

        # The prompt contains at most 400 Xs
        assert "X" * 400 in captured["prompt"]
        assert "X" * 401 not in captured["prompt"]


# ---------------------------------------------------------------------------
# generate_h1_draft — error handling
# ---------------------------------------------------------------------------
class TestGenerateH1DraftErrors:
    async def test_returns_none_on_llm_exception(self):
        with patch("sentinel.services.ai_explain.call_llm", new=AsyncMock(side_effect=RuntimeError("LLM down"))):
            result = await generate_h1_draft(_finding(), scan=None, target="test")
        assert result is None

    async def test_returns_none_on_network_error(self):
        import httpx
        with patch("sentinel.services.ai_explain.call_llm", new=AsyncMock(side_effect=httpx.ConnectError("timeout"))):
            result = await generate_h1_draft(_finding(), scan=None, target="test")
        assert result is None

    async def test_works_with_orm_like_object(self):
        """Finding can be an ORM object (has attributes), not just a dict."""
        class FakeFinding:
            tool = "bandit"
            rule = "B301"
            severity = "ERROR"
            confidence = "likely"
            file = "model.py"
            line = 99
            message = "pickle.load() called"
            owasp = "A08"
            remediation = "Use safetensors"

        with patch("sentinel.services.ai_explain.call_llm", new=AsyncMock(return_value=_MOCK_DRAFT)):
            result = await generate_h1_draft(FakeFinding(), scan=None, target="test-space")
        assert result == _MOCK_DRAFT

    async def test_works_with_none_scan(self):
        """scan=None should not crash."""
        with patch("sentinel.services.ai_explain.call_llm", new=AsyncMock(return_value=_MOCK_DRAFT)):
            result = await generate_h1_draft(_finding(), scan=None, target="test")
        assert result is not None