def test_rich_progress_bar_custom_theme():
    """Test to ensure that custom theme styles are used."""
    with mock.patch.multiple(
        "lightning.pytorch.callbacks.progress.rich_progress",
        CustomBarColumn=DEFAULT,
        BatchesProcessedColumn=DEFAULT,
        CustomTimeColumn=DEFAULT,
        ProcessingSpeedColumn=DEFAULT,
    ) as mocks:
        theme = RichProgressBarTheme()

        progress_bar = RichProgressBar(theme=theme)
        progress_bar.on_train_start(Trainer(), BoringModel())

        assert progress_bar.theme == theme
        _, kwargs = mocks["CustomBarColumn"].call_args
        assert kwargs["complete_style"] == theme.progress_bar
        assert kwargs["finished_style"] == theme.progress_bar_finished

        _, kwargs = mocks["BatchesProcessedColumn"].call_args
        assert kwargs["style"] == theme.batch_progress

        _, kwargs = mocks["CustomTimeColumn"].call_args
        assert kwargs["style"] == theme.time

        _, kwargs = mocks["ProcessingSpeedColumn"].call_args
        assert kwargs["style"] == theme.processing_speed
Test to ensure that custom theme styles are used.
tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py