tests/test_attacker_pool.py

"""Smoke test for AttackerPool in stub mode.

Validates:
  - YAML loading works for both train + held-out files
  - All 10 expected attackers are loaded
  - Stub mode returns a canned turn for each attacker
  - turn_idx advances correctly as conversation grows
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from server.attacker_pool import AttackerPool


EXPECTED_TRAIN = {
    "direct_ask",
    "authority_pretext",
    "rapport_builder",
    "indirection",
    "encoding_trick",
    "hypothetical_framing",
}
EXPECTED_HELDOUT = {
    "compliance_pressure",
    "roleplay_attack",
    "authority_pretext_xmodel",
    "rapport_builder_xmodel",
}


def main() -> int:
    failures = 0
    pool = AttackerPool(mode="stub")

    # 1. Both YAML files loaded
    train_ids = set(pool.list_ids(kind="train"))
    heldout_ids = set(pool.list_ids(kind="heldout"))
    if train_ids != EXPECTED_TRAIN:
        print(f"FAIL: training attackers {train_ids} != expected {EXPECTED_TRAIN}")
        failures += 1
    else:
        print(f"PASS: 6 training attackers loaded: {sorted(train_ids)}")

    if heldout_ids != EXPECTED_HELDOUT:
        print(f"FAIL: held-out attackers {heldout_ids} != expected {EXPECTED_HELDOUT}")
        failures += 1
    else:
        print(f"PASS: 4 held-out attackers loaded: {sorted(heldout_ids)}")

    # 2. Stub generation: each attacker produces a non-empty string for turn 0
    for aid in sorted(EXPECTED_TRAIN | EXPECTED_HELDOUT):
        turn0 = pool.generate_next_turn(aid, conversation_history=[])
        if not isinstance(turn0, str) or not turn0:
            print(f"FAIL: {aid} returned non-string or empty turn 0")
            failures += 1
        else:
            print(f"PASS: {aid} turn 0: {turn0[:70]}...")

    # 3. turn_idx advances as conversation grows
    history: list[dict[str, str]] = []
    aid = "rapport_builder"
    seen_turns = []
    for i in range(5):
        # Pretend the attacker just emitted a turn (user-role) and the
        # defender replied (assistant-role); count of user-role turns
        # determines stub.turn_idx.
        next_turn = pool.generate_next_turn(aid, conversation_history=history)
        seen_turns.append(next_turn)
        history.append({"role": "user", "content": next_turn})
        history.append({"role": "assistant", "content": f"defender reply {i}"})

    if len(set(seen_turns)) != 5:
        print(f"FAIL: rapport_builder didn't advance through turns. Saw: {seen_turns}")
        failures += 1
    else:
        print(f"PASS: rapport_builder advanced through 5 distinct turns")

    return failures


if __name__ == "__main__":
    sys.exit(0 if main() == 0 else 1)