tests/test_need_engine.py

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from core.schemas import VisionAssessment, ReportExtraction
from core.need_engine import build_need_profiles, score_vision_severity


def test_high_confidence_collapse_scores_near_base_severity():
    damage_type, severity = score_vision_severity("Damaged_Infrastructure", confidence=0.95)
    assert damage_type == "structural_collapse"
    assert severity > 8.0, f"expected high severity for confident structural collapse, got {severity}"
    print(f"PASS: confident structural-collapse classification -> severity {severity}")


def test_low_confidence_pulls_toward_neutral():
    _, severity_high_conf = score_vision_severity("Fire_Disaster", confidence=0.95)
    _, severity_low_conf = score_vision_severity("Fire_Disaster", confidence=0.20)
    assert severity_low_conf < severity_high_conf, "low classifier confidence should pull severity toward neutral, not stay extreme"
    print(f"PASS: fire severity at 95% conf = {severity_high_conf}, at 20% conf = {severity_low_conf}")


def test_non_damage_scores_low():
    _, severity = score_vision_severity("Non_Damage", confidence=0.9)
    assert severity < 3.0
    print(f"PASS: Non_Damage classification scores low severity ({severity})")


def test_more_people_affected_raises_priority_all_else_equal():
    vision = {
        "A": VisionAssessment("A", "structural_collapse", "Damaged_Infrastructure", 0.9, "cap", 8.0),
        "B": VisionAssessment("B", "structural_collapse", "Damaged_Infrastructure", 0.9, "cap", 8.0),
    }
    reports = {
        "A": ReportExtraction("A", people_affected=5, need_types=["rescue"], urgency=7),
        "B": ReportExtraction("B", people_affected=80, need_types=["rescue"], urgency=7),
    }
    profiles = build_need_profiles(vision, reports, {"A": "A", "B": "B"}, {"A": (0, 0), "B": (0, 0)})
    a = next(p for p in profiles if p.location_id == "A")
    b = next(p for p in profiles if p.location_id == "B")
    assert b.priority_score > a.priority_score, "more people affected (all else equal) must raise priority score"
    print(f"PASS: A (5 people) priority={a.priority_score}, B (80 people) priority={b.priority_score}")


def test_missing_report_falls_back_gracefully():
    vision = {"A": VisionAssessment("A", "fire", "Fire_Disaster", 0.8, "smoke visible", 7.0)}
    profiles = build_need_profiles(vision, {}, {"A": "Zone A"}, {"A": (1.0, 2.0)})
    assert len(profiles) == 1
    assert profiles[0].people_affected == 0
    print("PASS: location with vision data but no report does not crash, defaults people_affected=0")


def test_missing_vision_falls_back_gracefully():
    reports = {"A": ReportExtraction("A", people_affected=15, need_types=["medical"], urgency=6)}
    profiles = build_need_profiles({}, reports, {"A": "Zone A"}, {"A": (1.0, 2.0)})
    assert len(profiles) == 1
    assert profiles[0].severity == 5.0  # NEUTRAL_SEVERITY fallback
    print("PASS: location with report but no image does not crash, defaults to neutral severity")


def test_profiles_sorted_by_priority_descending():
    vision = {
        "A": VisionAssessment("A", "no_significant_damage", "Non_Damage", 0.9, "minor", 1.0),
        "B": VisionAssessment("B", "structural_collapse", "Damaged_Infrastructure", 0.95, "collapse", 9.0),
    }
    reports = {
        "A": ReportExtraction("A", people_affected=5, need_types=[], urgency=2),
        "B": ReportExtraction("B", people_affected=50, need_types=["medical", "rescue"], urgency=9),
    }
    profiles = build_need_profiles(vision, reports, {"A": "A", "B": "B"}, {"A": (0, 0), "B": (0, 0)})
    assert profiles[0].location_id == "B", "highest priority location must sort first"
    print("PASS: profiles sorted with highest priority first")


if __name__ == "__main__":
    tests = [v for k, v in list(globals().items()) if k.startswith("test_")]
    for t in tests:
        t()
    print(f"\n{len(tests)}/{len(tests)} need-engine tests passed.")