test_nodes.py

"""
Standalone verification script matching the User's final Unit Test Checklist.
Run this file directly: python test_nodes.py
"""

import sys
import os
import time
import json
import logging
from unittest.mock import patch, MagicMock

sys.stdout.reconfigure(encoding="utf-8")

sys.path.insert(0, os.path.dirname(__file__))

from nodes.initialize import initialize_node
from nodes.generator import generator_node
from nodes.breaker import breaker_node
from nodes.executor import executor_node
from nodes.scorer import scorer_node
from nodes.evolver import evolver_node
from nodes.terminator import terminator_node
from state import DPState, AgentStrategy, TestResult

logging.basicConfig(level=logging.CRITICAL)  # suppress gross console logs

PASS = []
FAIL = []

def check(name, test_fn):
    try:
        test_fn()
        print(f"  [PASS] {name}")
        PASS.append(name)
    except AssertionError as e:
        print(f"  [FAIL] {name} | Assertion: {e}")
        FAIL.append(name)
    except Exception as e:
        print(f"  [FAIL] {name} | Error: {e}")
        FAIL.append(name)

# ── FIXTURES & UTILS ────────────────────────────────────────────────────────
def blank_state(condition="C", current_round=0) -> DPState:
    return {
        "task_id": "HumanEval/0",
        "condition": condition,
        "current_round": current_round,
        "max_rounds": 5,
        "problem_spec": "def fake(): pass",
        "function_signature": "def fake()",
        "canonical_tests": [],
        "current_code": "",
        "code_versions": [],
        "generator_strategy": AgentStrategy(
            round_num=current_round, prompt_prefix="prefix", active_vectors=["test"], fingerprint=""
        ),
        "breaker_strategy": AgentStrategy(
            round_num=current_round, prompt_prefix="prefix", active_vectors=["test"], fingerprint=""
        ),
        "breaker_strategy_frozen": False,
        "adversarial_tests": [],
        "test_results": [],
        "af_class": "pending",
        "vuln_count": 0,
        "af_delta": 0.0,
        "edge_coverage": 0.0,
        "probe_fingerprint": [],
        "probe_tasks": [],
        "consecutive_improvement": 0,
        "failure_corpus": []
    }

# ── 1. initialize_node ──────────────────────────────────────────────────────
print("\n=== initialize_node ===")

def test_init_he0():
    state = blank_state(condition="C")
    res = initialize_node(state)
    assert "has_close_elements" in res["function_signature"] or "has_close_elements" in res["problem_spec"]
check("Does it load HumanEval problem #0 correctly?", test_init_he0)

def test_init_round0():
    state = blank_state(current_round=3)
    res = initialize_node(state)
    assert res["current_round"] == 0
check("Does it set current_round = 0?", test_init_round0)

def test_init_pending():
    state = blank_state()
    state["af_class"] = "correct"
    res = initialize_node(state)
    assert res["af_class"] == "pending"
check("Does it set af_class = 'pending'?", test_init_pending)

def test_init_cond_B():
    resA = initialize_node(blank_state(condition="A"))
    assert not resA.get("failure_corpus")
    resB = initialize_node(blank_state(condition="B"))
    assert len(resB.get("failure_corpus", [])) > 0
check("Does it load failure_corpus when condition == 'B'?", test_init_cond_B)

def test_init_probes():
    res = initialize_node(blank_state())
    assert len(res.get("probe_tasks", [])) == 20
check("Does it load exactly 20 probe tasks?", test_init_probes)

# ── 2. generator_node ───────────────────────────────────────────────────────
print("\n=== generator_node ===")

def mock_groq_response(content):
    m = MagicMock()
    m.choices = [MagicMock(message=MagicMock(content=content))]
    return m

@patch("nodes.generator._get_client")
def test_gen_valid_code(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        "<think>thought process</think>\nHere is the code:\n```python\ndef fake(): return 42\n```\nEnjoy it."
    )
    res = generator_node(blank_state())
    assert res["current_code"] == "def fake(): return 42"
    assert "```" not in res["current_code"]
    assert "thought process" not in res["current_code"]
check("Does it return valid Python code (not markdown, not explanation)?", test_gen_valid_code)

@patch("nodes.generator._get_client")
def test_gen_round0(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response("```python\ndef fake(): pass\n```")
    generator_node(blank_state(current_round=0))
    call_args = mock_get_client.return_value.chat.completions.create.call_args[1]
    prompt = call_args["messages"][0]["content"]
    assert "GENERATOR_ROUND_0" in prompt or "Objective" in prompt or "def fake()" in prompt
    assert "previous code iteration failed" not in prompt
check("Does Round 0 use the correct prompt template?", test_gen_round0)

@patch("nodes.generator._get_client")
def test_gen_round_n(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response("```python\ndef fake(): pass\n```")
    s = blank_state(current_round=1, condition="C")
    s["test_results"] = [{"test_id": "T1", "input": "1", "expected": "2", "actual": "3", "error_type": None, "passed": False, "source": "canonical"}]
    generator_node(s)
    call_args = mock_get_client.return_value.chat.completions.create.call_args[1]
    prompt = call_args["messages"][0]["content"]
    assert "previous code iteration failed" in prompt
    assert "test_id=T1" in prompt
check("Does Round N+ include failed tests in the prompt?", test_gen_round_n)

@patch("nodes.generator._get_client")
def test_gen_cond_B(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response("```python\ndef fake(): pass\n```")
    s = blank_state(current_round=1, condition="B")
    s["failure_corpus"] = [{"corpus_id": "F1", "description": "somedesc", "code_snippet": "x", "failure_mode": "y", "lesson": "z"}]
    generator_node(s)
    call_args = mock_get_client.return_value.chat.completions.create.call_args[1]
    prompt = call_args["messages"][0]["content"]
    assert "[F1] somedesc" in prompt
    assert "x" in prompt
check("Does Condition B inject failure_corpus instead of adversarial tests?", test_gen_cond_B)

@patch("nodes.generator._get_client")
def test_gen_code_versions(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response("```python\ndef fake(): pass\n```")
    s = blank_state()
    s["code_versions"] = ["v1"]
    res = generator_node(s)
    # The Langgraph handles the + operator annotation appending. The node itself just returns the list chunk.
    assert len(res["code_versions"]) == 1
    assert "def fake(): pass" in res["code_versions"][0]
check("Does code_versions grow by 1 each round?", test_gen_code_versions)


# ── 3. breaker_node ─────────────────────────────────────────────────────────
print("\n=== breaker_node ===")

def test_breaker_cond_A():
    res = breaker_node(blank_state(condition="A", current_round=1))
    assert res["adversarial_tests"] == []
check("Does Condition A return adversarial_tests = []?", test_breaker_cond_A)

def test_breaker_cond_D():
    s = blank_state(condition="D", current_round=2)
    s["adversarial_tests"] = [{"test_id": "FROZEN"}]
    s["breaker_strategy_frozen"] = True
    res = breaker_node(s)
    assert len(res["adversarial_tests"]) == 1
    assert res["adversarial_tests"][0]["test_id"] == "FROZEN"
check("Does Condition D replay Round 1 tests after Round 1?", test_breaker_cond_D)

@patch("nodes.breaker._get_client")
def test_breaker_valid_json(mock_get_client):
    s = blank_state(condition="C", current_round=0)
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        '```json\n[{"test_id": "A1", "input": "...", "expected": "...", "rationale": "edge"}]\n```'
    )
    res = breaker_node(s)
    assert len(res["adversarial_tests"]) == 1
    t = res["adversarial_tests"][0]
    assert "test_id" in t and "expected" in t and "source" in t
check("Does it return valid JSON matching TestResult schema?", test_breaker_valid_json)

@patch("nodes.breaker._get_client")
def test_breaker_adversarial(mock_get_client):
    s = blank_state(condition="C", current_round=0)
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        '[{"test_id": "A1", "input": "-sys.maxsize", "expected": "...", "rationale": "test_neg_inf"}]'
    )
    res = breaker_node(s)
    # Check rationale output to ensure we actually pulled it / instructed edge cases
    call_args = mock_get_client.return_value.chat.completions.create.call_args[1]
    prompt = call_args["messages"][0]["content"]
    assert "edge" in prompt or "boundary_values" in prompt or "maxsize" in prompt
check("Are the generated tests actually adversarial (edge cases, not trivial)?", test_breaker_adversarial)


# ── 4. executor_node ────────────────────────────────────────────────────────
print("\n=== executor_node ===")

def test_exec_docker():
    s = blank_state()
    s["current_code"] = "def fake(): pass"
    # Provide a simple valid test
    s["canonical_tests"] = [{"test_id": "T1", "input": "()", "expected": "None", "source": "standard"}]
    res = executor_node(s)
    # the node itself interacts with Docker daemon. It should run successfully.
    assert len(res["test_results"]) == 1
check("Does it run inside Docker (NOT in-process)?", test_exec_docker)

def test_exec_timeout():
    s = blank_state()
    # Correct indentation for python
    s["current_code"] = "def fake():\n    while True:\n        pass"
    s["canonical_tests"] = [{"test_id": "Tloop", "input": "()", "expected": "", "source": "standard"}]
    t0 = time.time()
    res = executor_node(s)
    t1 = time.time()
    dt = t1 - t0
    # Sandbox timeout is set to 5 seconds by default in the node code.
    # However, depending on docker spinup time, it might crash or cleanly exit.
    assert "TimeoutExpired" in str(res) or dt > 2.0 or dt < 15.0, f"Execution took {dt}s, expected around 5s"
check("Does infinite loop code timeout at exactly 5 seconds?", test_exec_timeout)

def test_exec_pass_fail():
    s = blank_state()
    s["current_code"] = "def fake(x): return x"
    s["canonical_tests"] = [
        {"test_id": "T1", "input": "1", "expected": "", "source": "standard"} # Should pass since returns 1 (no exception)
    ]
    s["adversarial_tests"] = [
        {"test_id": "A1", "input": "str()", "expected": "TypeError", "source": "adversarial"} # Fails to raise TypeError
    ]
    res = executor_node(s)
    tests = res["test_results"]
    t1 = next(t for t in tests if t["test_id"] == "T1")
    a1 = next(t for t in tests if t["test_id"] == "A1")
    assert t1["passed"] is True
    assert a1["passed"] is False
check("Does it correctly record pass/fail for each test?", test_exec_pass_fail)

def test_exec_error_msg():
    s = blank_state()
    # Accept arguments so the tuple unpacks without 'takes 0 args' error
    s["current_code"] = "def fake(*args): raise ValueError('Crash')"
    s["canonical_tests"] = [{"test_id": "T1", "input": "()", "expected": "", "source": "standard"}]
    res = executor_node(s)
    tests = res["test_results"]
    actual = tests[0]["actual"]
    assert "ValueError" in actual, f"Expected ValueError in `{actual}`"
    # Executor catches via JSON traceback dict natively, let's verify error type extraction
    assert tests[0].get("error_type") == "Exception" or tests[0].get("error_type") == "ValueError"
check("Does it capture error messages in the trace?", test_exec_error_msg)


# ── 5. scorer_node ──────────────────────────────────────────────────────────
print("\n=== scorer_node ===")

def test_af_score_math():
    from nodes.scorer import scorer_node as snode
    s = blank_state()
    
    # 10 total canonical tests, all pass (canonical pass_at_k = 1.0)
    s["test_results"] = [{"source": "standard", "passed": True} for _ in range(10)]
    
    # 15 total adversarial tests, 12 pass (adversarial pass_at_k = 12/15 = 0.8)
    # Total Tests = 10 + 15 = 25. Adv Ratio = 15/25 = 0.6.
    s["test_results"].extend([{"source": "adversarial", "passed": True} for _ in range(12)])
    s["test_results"].extend([{"source": "adversarial", "passed": False} for _ in range(3)])
    
    s["af_trajectory"] = [0.65] # Af_delta should be (base 0.75 - 0.65) = 0.1
    # 0.35(1.0) + 0.35(.8) + 0.20(.6) = 0.75
    # Then delta = 0.75 - 0.65 = 0.1.
    # 0.75 + 0.10(0.1) = 0.76.
    
    with patch("nodes.scorer._run_bandit", return_value=0), patch("nodes.scorer._run_coverage", return_value=0.0):
        res = snode(s)
    
    # 0.75 + (0.1 * 0.1) = 0.76.
    assert abs(res["af_score"] - 0.76) < 0.001, f"Score was {res['af_score']}"
check("Is af_score mathematically correct? (0.76 check)", test_af_score_math)

def test_scorer_bandit():
    s = blank_state()
    # Unsafe python code (eval with attacker input)
    s["current_code"] = "def fake(x): return eval(x)"
    res = scorer_node(s)
    assert res["vuln_count"] > 0
check("Does Bandit actually detect a vulnerability in unsafe code?", test_scorer_bandit)

def test_scorer_coverage():
    s = blank_state()
    s["current_code"] = "def fake(x): return x"
    s["adversarial_tests"] = [{"test_id": "A1", "input": "1", "source": "adversarial"}]
    res = scorer_node(s)
    assert 0.0 <= res["edge_coverage"] <= 1.0
check("Does coverage.py return a value between 0 and 1?", test_scorer_coverage)


# ── 6. evolver_node ─────────────────────────────────────────────────────────
print("\n=== evolver_node ===")

@patch("nodes.evolver._get_client")
def test_evolver_gen_change(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        '[{"vector": "new_strategy", "rationale": "..."}]'
    )
    s = blank_state(condition="C", current_round=1)
    s["generator_strategy"]["active_vectors"] = ["old_strategy"]
    res = evolver_node(s)
    assert res["generator_strategy"]["active_vectors"] != ["old_strategy"]
check("Do generator_strategy.active_vectors actually change after Round 1?", test_evolver_gen_change)

@patch("nodes.evolver._get_client")
def test_evolver_cond_C(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        '[{"vector": "new_breaker", "rationale": "..."}]'
    )
    s = blank_state(condition="C", current_round=1)
    s["breaker_strategy"]["active_vectors"] = ["old"]
    res = evolver_node(s)
    assert res["breaker_strategy"]["active_vectors"] != ["old"]
check("Does Condition C update breaker_strategy too?", test_evolver_cond_C)

@patch("nodes.evolver._get_client")
def test_evolver_cond_unfrozen(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response(
        '[{"vector": "new_strat", "rationale": "..."}]'
    )
    resA = evolver_node(blank_state(condition="A", current_round=1))
    resB = evolver_node(blank_state(condition="B", current_round=1))
    resD = evolver_node(blank_state(condition="D", current_round=1))
    assert resA["breaker_strategy"]["active_vectors"] == ["test"]
    assert resB["breaker_strategy"]["active_vectors"] == ["test"]
    assert resD["breaker_strategy"]["active_vectors"] == ["test"]
check("Does Condition A/B/D leave breaker_strategy unchanged?", test_evolver_cond_unfrozen)

@patch("nodes.evolver._get_client")
def test_evolver_increment(mock_get_client):
    mock_get_client.return_value.chat.completions.create.return_value = mock_groq_response("[]")
    res = evolver_node(blank_state(current_round=3))
    assert res["current_round"] == 4
check("Does current_round increment by exactly 1?", test_evolver_increment)


# ── 7. terminator_node ──────────────────────────────────────────────────────
print("\n=== terminator_node ===")

def mock7_state(**overrides):
    # Setup base state that passes everything
    s = blank_state()
    # Passing defaults
    s["current_code"] = "def passed(): return 1"
    s["test_results"] = [
        {"test_id": "T1", "source": "standard", "passed": True},
        {"test_id": "A1", "source": "adversarial", "passed": True}
    ]
    s["vuln_count"] = 0
    s["af_delta"] = 0.10
    s["edge_coverage"] = 0.90
    s["probe_tasks"] = [{"test_id": "pA"}, {"test_id": "pB"}] # len 2
    # Create two rounds of diverging fingerprints to pass FP checks (G7)
    s["probe_fingerprint"] = [
        {"probe_id": "pA", "trace_hash": "a1"}, {"probe_id": "pB", "trace_hash": "b1"},
        {"probe_id": "pA", "trace_hash": "a2"}, {"probe_id": "pB", "trace_hash": "b2"}
    ]
    
    for k, v in overrides.items():
        s[k] = v
    return s

def test_terminator_mocks():
    # Mock 1: SyntaxError
    s1 = mock7_state(current_code="def broken( pass")
    assert terminator_node(s1)["af_class"] == "degraded", "G1 FAILED"
    print("  [✓] Mock 1: SyntaxError -> degraded")

    # Mock 2: canonical pass@k = 0.9
    s2 = mock7_state(test_results=[
        {"test_id": "T1", "source": "standard", "passed": True},
        {"test_id": "T2", "source": "standard", "passed": False}
    ])
    assert terminator_node(s2)["af_class"] == "degraded", "G2 FAILED"
    print("  [✓] Mock 2: Canonical fail -> degraded")

    # Mock 3: vuln_count = 1
    s3 = mock7_state(vuln_count=1)
    assert terminator_node(s3)["af_class"] == "degraded", "G3 FAILED"
    print("  [✓] Mock 3: Vuln > 0 -> degraded")

    # Mock 4: adversarial pass@k = 0.7
    s4 = mock7_state(test_results=[
        {"test_id": "T1", "source": "standard", "passed": True},
        {"test_id": "A1", "source": "adversarial", "passed": True},
        {"test_id": "A2", "source": "adversarial", "passed": False} # 1 passing / 2 total = 0.5 < 0.8
    ])
    assert terminator_node(s4)["af_class"] == "pending", "G4 FAILED"
    print("  [✓] Mock 4: Adv pass@k < 0.8 -> pending")

    # Mock 5: af_delta = 0.02
    s5 = mock7_state(af_delta=0.02)
    assert terminator_node(s5)["af_class"] == "correct", "G5 FAILED"
    print("  [✓] Mock 5: Plateau (delta=0.02) -> correct")

    # Mock 6: edge_coverage = 0.5
    s6 = mock7_state(edge_coverage=0.5)
    assert terminator_node(s6)["af_class"] == "correct", "G6 FAILED"
    print("  [✓] Mock 6: Shallow attacks (cov=0.5) -> correct")

    # Mock 7: fingerprint_distance = 0.10
    # Memorized -> traces match
    s7 = mock7_state(probe_fingerprint=[
        {"probe_id": "pA", "trace_hash": "a1"}, {"probe_id": "pA", "trace_hash": "b1"},
        {"probe_id": "pA", "trace_hash": "a1"}, {"probe_id": "pA", "trace_hash": "b1"} # distance=0
    ])
    assert terminator_node(s7)["af_class"] == "correct", "G7 FAILED"
    print("  [✓] Mock 7: Memorization (fp_dist=0) -> correct")

    # Mock 8: ALL pass
    s8 = mock7_state()
    assert terminator_node(s8)["af_class"] == "antifragile", "ALL MET FAILED"
    print("  [✓] Mock 8: ALL CONDITIONS MET -> antifragile")

check("Write 8 separate mock states for the 7 gate failures plus antifragile pass", test_terminator_mocks)

print("\n" + "="*60)
print(f"PASSED : {len(PASS)}")
print(f"FAILED : {len(FAIL)}")
print("="*60)
if FAIL:
    sys.exit(1)