tests/test_drama.py

from copy import deepcopy

from drama import observed_index, train_and_score
from feed import fixture_choices


def event(event_id, date, home_id, away_id, home_score, away_score, state="post", stage="group-stage", cards=1):
    details = [
        {"type": {"text": "Yellow Card"}, "clock": {"value": 1800}}
        for _ in range(cards)
    ]
    details.append({"type": {"text": "Goal"}, "clock": {"value": 5100}})
    return {
        "id": str(event_id),
        "date": date,
        "season": {"slug": stage},
        "status": {
            "displayClock": "90'",
            "type": {"state": state, "completed": state == "post", "description": "Full Time"},
        },
        "competitions": [{
            "details": details,
            "competitors": [
                {
                    "homeAway": "home",
                    "score": str(home_score),
                    "team": {"id": str(home_id), "displayName": f"Team {home_id}"},
                    "statistics": [
                        {"name": "foulsCommitted", "displayValue": "12"},
                        {"name": "totalShots", "displayValue": "11"},
                    ],
                },
                {
                    "homeAway": "away",
                    "score": str(away_score),
                    "team": {"id": str(away_id), "displayName": f"Team {away_id}"},
                    "statistics": [
                        {"name": "foulsCommitted", "displayValue": "10"},
                        {"name": "totalShots", "displayValue": "9"},
                    ],
                },
            ],
        }],
    }


def tournament_sample(count=40):
    games = []
    for i in range(count):
        games.append(event(i, f"2026-06-{i % 28 + 1:02d}T12:00Z", i % 8, (i + 1) % 8, i % 4, (i + 2) % 3, cards=i % 5))
    return games


def test_cards_raise_observed_index():
    calm = event(1, "2026-06-01T12:00Z", 1, 2, 1, 0, cards=0)
    heated = deepcopy(calm)
    heated["competitions"][0]["details"] += [
        {"type": {"text": "Yellow Card"}, "clock": {"value": 4000}}
        for _ in range(5)
    ]

    assert observed_index(heated) > observed_index(calm)


def test_model_trains_on_completed_tournament_matches():
    games = tournament_sample()
    target = event(99, "2026-07-15T19:00Z", 1, 4, 0, 0, state="pre", stage="semifinals")
    result = train_and_score(games + [target], target)

    assert result["samples"] == 40
    assert 0 <= result["forecast"] <= 100
    assert len(result["importances"]) == 9


def test_fixture_choices_put_live_match_first():
    finished = event(1, "2026-07-14T19:00Z", 1, 2, 2, 1)
    live = event(2, "2026-07-15T19:00Z", 3, 4, 1, 1, state="in")
    upcoming = event(3, "2026-07-16T19:00Z", 5, 6, 0, 0, state="pre")

    choices, default = fixture_choices([finished, upcoming, live])

    assert default == "2"
    assert choices[0][1] == "2"