from app.dns_policy import collect_domain_dns_policy, parse_dmarc_records, parse_spf_records


def test_parse_dmarc_record_extracts_policy_tags():
    parsed, errors = parse_dmarc_records(["v=DMARC1; p=reject; sp=quarantine; pct=50; adkim=s; aspf=r; rua=mailto:d@example.com"])

    assert errors == []
    assert parsed.p == "reject"
    assert parsed.sp == "quarantine"
    assert parsed.pct == 50
    assert parsed.adkim == "s"
    assert parsed.rua == "mailto:d@example.com"


def test_parse_spf_record_extracts_includes_and_all_mechanism():
    parsed, errors = parse_spf_records(["v=spf1 include:_spf.google.com include:mailgun.org -all"])

    assert errors == []
    assert parsed.includes == ["_spf.google.com", "mailgun.org"]
    assert parsed.all_mechanism == "-all"


def test_collect_domain_dns_policy_queries_only_domain_dns_records():
    txt_records = {
        "_dmarc.example.com": ["v=DMARC1; p=reject; pct=100"],
        "example.com": ["v=spf1 include:_spf.example.net -all"],
    }
    queried = []

    def txt_lookup(name: str) -> list[str]:
        queried.append(name)
        if name not in txt_records:
            raise RuntimeError("not found")
        return txt_records[name]

    policy = collect_domain_dns_policy(
        "example.com",
        txt_lookup=txt_lookup,
        mx_lookup=lambda name: ["10 mail.example.com"],
    )

    assert policy.dmarc.p == "reject"
    assert policy.spf.all_mechanism == "-all"
    assert policy.mx_records == ["10 mail.example.com"]
    assert queried == ["_dmarc.example.com", "example.com"]
    assert policy.errors == []