def test_rich_progress_bar_reset_bars():
    """Test that the progress bar resets all internal bars when a new trainer stage begins."""
    bar = RichProgressBar()
    assert bar.is_enabled
    assert bar.progress is None
    assert bar._progress_stopped is False

    def _set_fake_bar_ids():
        bar.train_progress_bar_id = 0
        bar.val_sanity_progress_bar_id = 1
        bar.val_progress_bar_id = 2
        bar.test_progress_bar_id = 3
        bar.predict_progress_bar_id = 4

    for stage in ("train", "sanity_check", "validation", "test", "predict"):
        hook_name = f"on_{stage}_start"
        hook = getattr(bar, hook_name)

        _set_fake_bar_ids()  # pretend that bars are initialized from a previous run
        hook(Mock(), Mock())
        bar.teardown(Mock(), Mock(), Mock())

        # assert all bars are reset
        assert bar.train_progress_bar_id is None
        assert bar.val_sanity_progress_bar_id is None
        assert bar.val_progress_bar_id is None
        assert bar.test_progress_bar_id is None
        assert bar.predict_progress_bar_id is None

        # the progress object remains in case we need it for the next stage
        assert bar.progress is not None
Test that the progress bar resets all internal bars when a new trainer stage begins.
tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py