tests/test_extractor.py

"""
Test agent untuk memverifikasi kualitas ekstraksi JSON dari model.
Run: python tests/test_extractor.py
"""
import sys

sys.path.insert(0, ".")

from model.extractor import extract_entities

TEST_CASES = [
    {
        "id": "TC001",
        "name": "Standard narrative dream",
        "input": (
            "I was running through a forest at night. A white deer kept appearing. "
            "My childhood friend was there but their face kept changing."
        ),
        "required_fields": [
            "title",
            "mood",
            "characters",
            "places",
            "symbols",
            "connections",
            "core_theme",
        ],
        "min_characters": 1,
        "min_places": 1,
        "valid_moods": ["mysterious", "joyful", "anxious", "surreal", "peaceful"],
    },
    {
        "id": "TC002",
        "name": "Short dream",
        "input": "I was floating. That was it.",
        "required_fields": ["title", "mood", "core_theme"],
        "min_characters": 0,
        "min_places": 0,
        "valid_moods": ["mysterious", "joyful", "anxious", "surreal", "peaceful"],
    },
    {
        "id": "TC003",
        "name": "Emotionally complex dream",
        "input": (
            "My mother was there but she was also a stranger. We were in a house "
            "that kept expanding. Every room I opened had versions of me at different "
            "ages. I felt grief and wonder at the same time."
        ),
        "required_fields": [
            "title",
            "mood",
            "characters",
            "places",
            "connections",
            "core_theme",
        ],
        "min_characters": 1,
        "min_places": 1,
        "valid_moods": ["mysterious", "joyful", "anxious", "surreal", "peaceful"],
    },
]


def run_tests():
    passed, failed = 0, 0
    results = []

    for tc in TEST_CASES:
        print(f"\n[{tc['id']}] {tc['name']}")
        print(f"  Input: {tc['input'][:60]}...")

        try:
            result, raw, parse_ok, parse_error, _ = extract_entities(tc["input"])
            errors = []

            if not parse_ok:
                errors.append(f"Parse/quality failed: {parse_error}")

            for field in tc["required_fields"]:
                if field not in result:
                    errors.append(f"Missing field: {field}")

            if "mood" in result and result["mood"] not in tc["valid_moods"]:
                errors.append(f"Invalid mood: {result['mood']}")

            if len(result.get("characters", [])) < tc.get("min_characters", 0):
                errors.append(f"Too few characters: {len(result.get('characters', []))}")
            if len(result.get("places", [])) < tc.get("min_places", 0):
                errors.append(f"Too few places: {len(result.get('places', []))}")

            if errors:
                print(f"  ❌ FAILED: {errors}")
                failed += 1
            else:
                print(
                    f"  ✅ PASSED — mood={result['mood']}, "
                    f"chars={len(result.get('characters', []))}, "
                    f"places={len(result.get('places', []))}, "
                    f"symbols={len(result.get('symbols', []))}"
                )
                passed += 1

            results.append(
                {
                    "id": tc["id"],
                    "status": "pass" if not errors else "fail",
                    "output": result,
                    "raw_len": len(raw),
                    "errors": errors,
                }
            )

        except Exception as e:
            print(f"  💥 ERROR: {e}")
            failed += 1
            results.append({"id": tc["id"], "status": "error", "error": str(e)})

    print(f"\n{'=' * 50}")
    print(f"Results: {passed}/{passed + failed} passed")
    return passed, failed, results


if __name__ == "__main__":
    p, f, _ = run_tests()
    sys.exit(0 if f == 0 else 1)