refactor: move patch decorators outside test functions in callback tests

Author: littlebullGit

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -204,30 +204,30 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0):
     assert pbar.predict_progress_bar.leave
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_tqdm_progress_bar_fast_dev_run(tmp_path):
-    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
-        model = BoringModel()
+    model = BoringModel()
 
-        trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
+    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
 
-        trainer.fit(model)
+    trainer.fit(model)
 
-        pbar = trainer.progress_bar_callback
+    pbar = trainer.progress_bar_callback
 
-        assert pbar.val_progress_bar.n == 1
-        assert pbar.val_progress_bar.total == 1
+    assert pbar.val_progress_bar.n == 1
+    assert pbar.val_progress_bar.total == 1
 
-        # the train progress bar should display 1 batch
-        assert pbar.train_progress_bar.total == 1
-        assert pbar.train_progress_bar.n == 1
+    # the train progress bar should display 1 batch
+    assert pbar.train_progress_bar.total == 1
+    assert pbar.train_progress_bar.n == 1
 
-        trainer.validate(model)
+    trainer.validate(model)
 
-        # the validation progress bar should display 1 batch
-        assert pbar.val_progress_bar.total == 1
-        assert pbar.val_progress_bar.n == 1
+    # the validation progress bar should display 1 batch
+    assert pbar.val_progress_bar.total == 1
+    assert pbar.val_progress_bar.n == 1
 
-        trainer.test(model)
+    trainer.test(model)
 
     # the test progress bar should display 1 batch
     assert pbar.test_progress_bar.total == 1
@@ -325,17 +325,17 @@ def test_tqdm_progress_bar_default_value(tmp_path):
 
 
 @mock.patch.dict(os.environ, {"COLAB_GPU": "1"})
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_tqdm_progress_bar_value_on_colab(tmp_path):
     """Test that Trainer will override the default in Google COLAB."""
-    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
-        trainer = Trainer(default_root_dir=tmp_path)
-        assert trainer.progress_bar_callback.refresh_rate == 20
+    trainer = Trainer(default_root_dir=tmp_path)
+    assert trainer.progress_bar_callback.refresh_rate == 20
 
-        trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar())
-        assert trainer.progress_bar_callback.refresh_rate == 20
+    trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar())
+    assert trainer.progress_bar_callback.refresh_rate == 20
 
-        trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar(refresh_rate=19))
-        assert trainer.progress_bar_callback.refresh_rate == 19
+    trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar(refresh_rate=19))
+    assert trainer.progress_bar_callback.refresh_rate == 19
 
 
 @pytest.mark.parametrize(

tests/tests_pytorch/callbacks/test_callbacks.py

@@ -119,21 +119,21 @@ def load_state_dict(self, state_dict) -> None:
         self.state = state_dict["state"]
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_resume_callback_state_saved_by_type_stateful(tmp_path):
     """Test that a legacy checkpoint that didn't use a state key before can still be loaded, using
     state_dict/load_state_dict."""
-    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
-        model = BoringModel()
-        callback = OldStatefulCallback(state=111)
-        trainer = Trainer(default_root_dir=tmp_path, max_steps=1, callbacks=[callback])
-        trainer.fit(model)
-        ckpt_path = Path(trainer.checkpoint_callback.best_model_path)
-        assert ckpt_path.exists()
-
-        callback = OldStatefulCallback(state=222)
-        trainer = Trainer(default_root_dir=tmp_path, max_steps=2, callbacks=[callback])
-        trainer.fit(model, ckpt_path=ckpt_path)
-        assert callback.state == 111
+    model = BoringModel()
+    callback = OldStatefulCallback(state=111)
+    trainer = Trainer(default_root_dir=tmp_path, max_steps=1, callbacks=[callback])
+    trainer.fit(model)
+    ckpt_path = Path(trainer.checkpoint_callback.best_model_path)
+    assert ckpt_path.exists()
+
+    callback = OldStatefulCallback(state=222)
+    trainer = Trainer(default_root_dir=tmp_path, max_steps=2, callbacks=[callback])
+    trainer.fit(model, ckpt_path=ckpt_path)
+    assert callback.state == 111
 
 
 def test_resume_incomplete_callbacks_list_warning(tmp_path):

tests/tests_pytorch/trainer/connectors/test_callback_connector.py

@@ -35,73 +35,73 @@
 from lightning.pytorch.trainer.connectors.callback_connector import _CallbackConnector
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_checkpoint_callbacks_are_last(tmp_path):
     """Test that checkpoint callbacks always come last."""
-    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
-        checkpoint1 = ModelCheckpoint(tmp_path / "path1", filename="ckpt1", monitor="val_loss_c1")
-        checkpoint2 = ModelCheckpoint(tmp_path / "path2", filename="ckpt2", monitor="val_loss_c2")
-        early_stopping = EarlyStopping(monitor="foo")
-        lr_monitor = LearningRateMonitor()
-        model_summary = ModelSummary()
-        progress_bar = TQDMProgressBar()
-
-        # no model reference
-        trainer = Trainer(callbacks=[checkpoint1, progress_bar, lr_monitor, model_summary, checkpoint2])
-        assert trainer.callbacks == [
-            progress_bar,
-            lr_monitor,
-            model_summary,
-            checkpoint1,
-            checkpoint2,
-        ]
-
-        # no model callbacks
-        model = LightningModule()
-        model.configure_callbacks = lambda: []
-        trainer.strategy._lightning_module = model
-        cb_connector = _CallbackConnector(trainer)
-        cb_connector._attach_model_callbacks()
-        assert trainer.callbacks == [
-            progress_bar,
-            lr_monitor,
-            model_summary,
-            checkpoint1,
-            checkpoint2,
-        ]
-
-        # with model-specific callbacks that substitute ones in Trainer
-        model = LightningModule()
-        model.configure_callbacks = lambda: [checkpoint1, early_stopping, model_summary, checkpoint2]
-        trainer = Trainer(callbacks=[progress_bar, lr_monitor, ModelCheckpoint(tmp_path, filename="ckpt_trainer")])
-        trainer.strategy._lightning_module = model
-        cb_connector = _CallbackConnector(trainer)
-        cb_connector._attach_model_callbacks()
-        assert trainer.callbacks == [
-            progress_bar,
-            lr_monitor,
-            early_stopping,
-            model_summary,
-            checkpoint1,
-            checkpoint2,
-        ]
-
-        # with tuner-specific callbacks that substitute ones in Trainer
-        model = LightningModule()
-        batch_size_finder = BatchSizeFinder()
-        model.configure_callbacks = lambda: [checkpoint2, early_stopping, batch_size_finder, model_summary, checkpoint1]
-        trainer = Trainer(callbacks=[progress_bar, lr_monitor])
-        trainer.strategy._lightning_module = model
-        cb_connector = _CallbackConnector(trainer)
-        cb_connector._attach_model_callbacks()
-        assert trainer.callbacks == [
-            batch_size_finder,
-            progress_bar,
-            lr_monitor,
-            early_stopping,
-            model_summary,
-            checkpoint2,
-            checkpoint1,
-        ]
+    checkpoint1 = ModelCheckpoint(tmp_path / "path1", filename="ckpt1", monitor="val_loss_c1")
+    checkpoint2 = ModelCheckpoint(tmp_path / "path2", filename="ckpt2", monitor="val_loss_c2")
+    early_stopping = EarlyStopping(monitor="foo")
+    lr_monitor = LearningRateMonitor()
+    model_summary = ModelSummary()
+    progress_bar = TQDMProgressBar()
+
+    # no model reference
+    trainer = Trainer(callbacks=[checkpoint1, progress_bar, lr_monitor, model_summary, checkpoint2])
+    assert trainer.callbacks == [
+        progress_bar,
+        lr_monitor,
+        model_summary,
+        checkpoint1,
+        checkpoint2,
+    ]
+
+    # no model callbacks
+    model = LightningModule()
+    model.configure_callbacks = lambda: []
+    trainer.strategy._lightning_module = model
+    cb_connector = _CallbackConnector(trainer)
+    cb_connector._attach_model_callbacks()
+    assert trainer.callbacks == [
+        progress_bar,
+        lr_monitor,
+        model_summary,
+        checkpoint1,
+        checkpoint2,
+    ]
+
+    # with model-specific callbacks that substitute ones in Trainer
+    model = LightningModule()
+    model.configure_callbacks = lambda: [checkpoint1, early_stopping, model_summary, checkpoint2]
+    trainer = Trainer(callbacks=[progress_bar, lr_monitor, ModelCheckpoint(tmp_path, filename="ckpt_trainer")])
+    trainer.strategy._lightning_module = model
+    cb_connector = _CallbackConnector(trainer)
+    cb_connector._attach_model_callbacks()
+    assert trainer.callbacks == [
+        progress_bar,
+        lr_monitor,
+        early_stopping,
+        model_summary,
+        checkpoint1,
+        checkpoint2,
+    ]
+
+    # with tuner-specific callbacks that substitute ones in Trainer
+    model = LightningModule()
+    batch_size_finder = BatchSizeFinder()
+    model.configure_callbacks = lambda: [checkpoint2, early_stopping, batch_size_finder, model_summary, checkpoint1]
+    trainer = Trainer(callbacks=[progress_bar, lr_monitor])
+    trainer.strategy._lightning_module = model
+    cb_connector = _CallbackConnector(trainer)
+    cb_connector._attach_model_callbacks()
+    assert trainer.callbacks == [
+        batch_size_finder,
+        progress_bar,
+        lr_monitor,
+        early_stopping,
+        model_summary,
+        checkpoint2,
+        checkpoint1,
+    ]
 
 
 class StatefulCallback0(Callback):
@@ -162,81 +162,81 @@ def test_all_callback_states_saved_before_checkpoint_callback(tmp_path):
     )
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_attach_model_callbacks():
     """Test that the callbacks defined in the model and through Trainer get merged correctly."""
-    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
-
-        def _attach_callbacks(trainer_callbacks, model_callbacks):
-            model = LightningModule()
-            model.configure_callbacks = lambda: model_callbacks
-            has_progress_bar = any(isinstance(cb, ProgressBar) for cb in trainer_callbacks + model_callbacks)
-            trainer = Trainer(
-                enable_checkpointing=False,
-                enable_progress_bar=has_progress_bar,
-                enable_model_summary=False,
-                callbacks=trainer_callbacks,
-            )
-            trainer.strategy._lightning_module = model
-            cb_connector = _CallbackConnector(trainer)
-            cb_connector._attach_model_callbacks()
-            return trainer
-
-        early_stopping1 = EarlyStopping(monitor="red")
-        early_stopping2 = EarlyStopping(monitor="blue")
-        progress_bar = TQDMProgressBar()
-        lr_monitor = LearningRateMonitor()
-        grad_accumulation = GradientAccumulationScheduler({1: 1})
-
-        # no callbacks
-        trainer = _attach_callbacks(trainer_callbacks=[], model_callbacks=[])
-        assert trainer.callbacks == []
-
-        # callbacks of different types
-        trainer = _attach_callbacks(trainer_callbacks=[early_stopping1], model_callbacks=[progress_bar])
-        assert trainer.callbacks == [early_stopping1, progress_bar]
-
-        # same callback type twice, different instance
-        trainer = _attach_callbacks(
-            trainer_callbacks=[progress_bar, EarlyStopping(monitor="red")],
-            model_callbacks=[early_stopping1],
-        )
-        assert trainer.callbacks == [progress_bar, early_stopping1]
-
-        # multiple callbacks of the same type in trainer
-        trainer = _attach_callbacks(
-            trainer_callbacks=[
-                LearningRateMonitor(),
-                EarlyStopping(monitor="yellow"),
-                LearningRateMonitor(),
-                EarlyStopping(monitor="black"),
-            ],
-            model_callbacks=[early_stopping1, lr_monitor],
-        )
-        assert trainer.callbacks == [early_stopping1, lr_monitor]
-
-        # multiple callbacks of the same type, in both trainer and model
-        trainer = _attach_callbacks(
-            trainer_callbacks=[
-                LearningRateMonitor(),
-                progress_bar,
-                EarlyStopping(monitor="yellow"),
-                LearningRateMonitor(),
-                EarlyStopping(monitor="black"),
-            ],
-            model_callbacks=[early_stopping1, lr_monitor, grad_accumulation, early_stopping2],
+
+    def _attach_callbacks(trainer_callbacks, model_callbacks):
+        model = LightningModule()
+        model.configure_callbacks = lambda: model_callbacks
+        has_progress_bar = any(isinstance(cb, ProgressBar) for cb in trainer_callbacks + model_callbacks)
+        trainer = Trainer(
+            enable_checkpointing=False,
+            enable_progress_bar=has_progress_bar,
+            enable_model_summary=False,
+            callbacks=trainer_callbacks,
         )
-        assert trainer.callbacks == [progress_bar, early_stopping1, lr_monitor, grad_accumulation, early_stopping2]
+        trainer.strategy._lightning_module = model
+        cb_connector = _CallbackConnector(trainer)
+        cb_connector._attach_model_callbacks()
+        return trainer
+
+    early_stopping1 = EarlyStopping(monitor="red")
+    early_stopping2 = EarlyStopping(monitor="blue")
+    progress_bar = TQDMProgressBar()
+    lr_monitor = LearningRateMonitor()
+    grad_accumulation = GradientAccumulationScheduler({1: 1})
+
+    # no callbacks
+    trainer = _attach_callbacks(trainer_callbacks=[], model_callbacks=[])
+    assert trainer.callbacks == []
+
+    # callbacks of different types
+    trainer = _attach_callbacks(trainer_callbacks=[early_stopping1], model_callbacks=[progress_bar])
+    assert trainer.callbacks == [early_stopping1, progress_bar]
+
+    # same callback type twice, different instance
+    trainer = _attach_callbacks(
+        trainer_callbacks=[progress_bar, EarlyStopping(monitor="red")],
+        model_callbacks=[early_stopping1],
+    )
+    assert trainer.callbacks == [progress_bar, early_stopping1]
+
+    # multiple callbacks of the same type in trainer
+    trainer = _attach_callbacks(
+        trainer_callbacks=[
+            LearningRateMonitor(),
+            EarlyStopping(monitor="yellow"),
+            LearningRateMonitor(),
+            EarlyStopping(monitor="black"),
+        ],
+        model_callbacks=[early_stopping1, lr_monitor],
+    )
+    assert trainer.callbacks == [early_stopping1, lr_monitor]
+
+    # multiple callbacks of the same type, in both trainer and model
+    trainer = _attach_callbacks(
+        trainer_callbacks=[
+            LearningRateMonitor(),
+            progress_bar,
+            EarlyStopping(monitor="yellow"),
+            LearningRateMonitor(),
+            EarlyStopping(monitor="black"),
+        ],
+        model_callbacks=[early_stopping1, lr_monitor, grad_accumulation, early_stopping2],
+    )
+    assert trainer.callbacks == [progress_bar, early_stopping1, lr_monitor, grad_accumulation, early_stopping2]
 
-        class CustomProgressBar(TQDMProgressBar): ...
+    class CustomProgressBar(TQDMProgressBar): ...
 
-        custom_progress_bar = CustomProgressBar()
-        # a custom callback that overrides ours
-        trainer = _attach_callbacks(trainer_callbacks=[progress_bar], model_callbacks=[custom_progress_bar])
-        assert trainer.callbacks == [custom_progress_bar]
+    custom_progress_bar = CustomProgressBar()
+    # a custom callback that overrides ours
+    trainer = _attach_callbacks(trainer_callbacks=[progress_bar], model_callbacks=[custom_progress_bar])
+    assert trainer.callbacks == [custom_progress_bar]
 
-        # edge case
-        bare_callback = Callback()
-        trainer = _attach_callbacks(trainer_callbacks=[bare_callback], model_callbacks=[custom_progress_bar])
+    # edge case
+    bare_callback = Callback()
+    trainer = _attach_callbacks(trainer_callbacks=[bare_callback], model_callbacks=[custom_progress_bar])
     assert trainer.callbacks == [bare_callback, custom_progress_bar]