cosmos-framework/cosmos_framework/utils/timer_test.py

import pytest

from cosmos_framework.utils.timer import Timer


def test_timer_decorator_resets_after_exception():
    timer = Timer(measure_cpu=True, measure_cuda=False, debug=True)
    calls = 0

    @timer
    def fail_once():
        nonlocal calls
        calls += 1
        if calls == 1:
            raise RuntimeError("first call failed")
        return "ok"

    with pytest.raises(RuntimeError, match="first call failed"):
        fail_once()

    assert timer.busy is False
    assert fail_once() == "ok"