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