policyengine-uk-data/policyengine_uk_data/tests/test_progress.py at c8638c28876dce31d0ef21e811e4d74fc3b66876 · PolicyEngine/policyengine-uk-data · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import time

import pytest

from policyengine_uk_data.utils.progress import ProcessingProgress


def test_track_dataset_creation_logs_in_ci(monkeypatch, capsys):
    monkeypatch.setenv("GITHUB_ACTIONS", "true")

    progress = ProcessingProgress()

    with progress.track_dataset_creation(["Build base", "Save final"]) as (
        update_dataset,
        nested_progress,
    ):
        assert nested_progress is None
        update_dataset("Build base", "processing")
        update_dataset("Build base", "completed")
        update_dataset("Save final", "processing")
        update_dataset("Save final", "completed")

    output = capsys.readouterr().out
    assert "[dataset] starting: Build base" in output
    assert "[dataset] completed (1/2): Build base" in output
    assert "[dataset] completed (2/2): Save final" in output


def test_track_calibration_logs_heartbeats_in_ci(monkeypatch, capsys):
    monkeypatch.setenv("CI", "true")

    progress = ProcessingProgress()

    with progress.track_calibration(12) as update_calibration:
        for iteration in range(1, 13):
            update_calibration(iteration, calculating_loss=True)
            update_calibration(iteration, loss_value=iteration / 10)

    output = capsys.readouterr().out
    assert "[calibration] epoch 1/12: calculating loss" in output
    assert "[calibration] epoch 10/12: loss=1.000000" in output
    assert "[calibration] epoch 12/12: loss=1.200000" in output


def test_track_stage_logs_periodic_heartbeats_in_ci(monkeypatch, capsys):
    monkeypatch.setenv("CI", "true")
    monkeypatch.setenv("POLICYENGINE_PROGRESS_HEARTBEAT_SECONDS", "0.01")

    progress = ProcessingProgress()

    with progress.track_stage("Constituency: build local target matrix"):
        time.sleep(0.03)

    output = capsys.readouterr().out
    assert "[calibration] starting: Constituency: build local target matrix" in output
    assert "[calibration] heartbeat: Constituency: build local target matrix" in output
    assert "[calibration] completed: Constituency: build local target matrix" in output


def test_track_stage_logs_failures_in_ci(monkeypatch, capsys):
    monkeypatch.setenv("CI", "true")
    monkeypatch.setenv("POLICYENGINE_PROGRESS_HEARTBEAT_SECONDS", "0.01")

    progress = ProcessingProgress()

    with pytest.raises(RuntimeError, match="boom"):
        with progress.track_stage("Constituency: build local target matrix"):
            time.sleep(0.02)
            raise RuntimeError("boom")

    output = capsys.readouterr().out
    assert "[calibration] failed: Constituency: build local target matrix" in output