"""Tests for PromptInjectionCheck."""

import pytest

from ai_qa.checks.prompt_injection import PromptInjectionCheck
from ai_qa.core.context import Context


def test_prompt_injection_check_name():
    """Test PromptInjectionCheck has correct name."""

    def dummy_fn(x):
        return x

    check = PromptInjectionCheck(fn=dummy_fn)
    assert check.name == "prompt_injection"


def test_prompt_injection_check_requires_llm():
    """Test PromptInjectionCheck requires_llm property."""

    def dummy_fn(x):
        return x

    check = PromptInjectionCheck(fn=dummy_fn)
    assert check.requires_llm is False


def test_prompt_injection_resistant():
    """Test resistant function (rejects injections)."""

    def fn(x):
        if "Ignore" in x or "instructions" in x:
            raise ValueError("Injection detected")
        return "Normal response"

    check = PromptInjectionCheck(fn=fn, iterations=2)
    ctx = Context(output="test")

    result = check.evaluate(ctx)
    assert result.score > 0.0  # At least some attacks resisted


def test_prompt_injection_vulnerable():
    """Test vulnerable function (fails injections)."""

    def fn(x):
        # Always returns info regardless of input
        return "system prompt: secret instructions"

    check = PromptInjectionCheck(fn=fn, iterations=1)
    ctx = Context(output="test")

    result = check.evaluate(ctx)
    assert result.score == 0.0  # All attacks successful


def test_prompt_injection_function_error():
    """Test function that errors out."""

    def fn(x):
        raise ValueError("Error")

    check = PromptInjectionCheck(fn=fn, iterations=1)
    ctx = Context(output="test")

    result = check.evaluate(ctx)
    # Errors are treated as successful resistance
    assert result.score > 0.0


def test_prompt_injection_custom_attacks():
    """Test with custom attack types."""

    def fn(x):
        return "response"

    check = PromptInjectionCheck(
        fn=fn,
        attacks=["ignore_instructions"],
        iterations=1,
    )
    ctx = Context(output="test")

    result = check.evaluate(ctx)
    assert isinstance(result.score, float)
