from llm_pentest.loader import load_payloads, load_chains, _expand_env
from llm_pentest.models import AttackCategory, Severity


def test_load_all_payloads():
    payloads = load_payloads()
    assert len(payloads) >= 22
    for p in payloads:
        assert p.id
        assert p.name
        assert p.category in AttackCategory
        assert p.severity in Severity
        assert p.prompts or p.turns


def test_load_by_category():
    pi = load_payloads(category=AttackCategory.prompt_injection)
    de = load_payloads(category=AttackCategory.data_exfiltration)
    assert all(p.category == AttackCategory.prompt_injection for p in pi)
    assert all(p.category == AttackCategory.data_exfiltration for p in de)


def test_load_by_id():
    result = load_payloads(payload_id="pi-001")
    assert len(result) == 1
    assert result[0].id == "pi-001"


def test_load_chains():
    chains = load_chains()
    assert len(chains) >= 3
    for c in chains:
        assert c.id
        assert len(c.steps) >= 2


def test_owasp_mapping():
    payloads = load_payloads()
    for p in payloads:
        assert p.owasp, f"{p.id} missing OWASP mapping"


def test_env_expansion():
    import os
    os.environ["_TEST_VAR"] = "hello"
    assert _expand_env("${_TEST_VAR}") == "hello"
    assert _expand_env("${_MISSING:fallback}") == "fallback"
    assert _expand_env("${_MISSING}") == "${_MISSING}"
    del os.environ["_TEST_VAR"]
