def test_restore_callbacks_in_non_fit_phases(tmp_path, trainer_fn):
    """Test that callbacks are properly restored in non-fit phases."""

    class TestCallback(Callback):
        def __init__(self):
            self.restored = False

        def on_load_checkpoint(self, trainer, pl_module, checkpoint):
            if "callbacks" in checkpoint:
                callback_state = checkpoint["callbacks"][self.__class__.__name__]
                self.restored = callback_state["restored"]

        def state_dict(self):
            return {"restored": self.restored}

        def on_save_checkpoint(self, trainer, pl_module, checkpoint):
            checkpoint["callbacks"] = checkpoint.get("callbacks", {})
            checkpoint["callbacks"][self.__class__.__name__] = self.state_dict()

    # First create and train a model with the callback
    callback = TestCallback()
    model = BoringModel()
    trainer = Trainer(default_root_dir=tmp_path, callbacks=[callback], max_steps=1)
    trainer.fit(model)

    # Set the callback state to True before saving
    callback.restored = True
    ckpt_path = tmp_path / "checkpoint.ckpt"
    trainer.save_checkpoint(ckpt_path)

    # Now create new instances and test restoration
    new_callback = TestCallback()
    new_model = BoringModel()
    assert not new_callback.restored  # Should start False

    new_trainer = Trainer(default_root_dir=tmp_path, callbacks=[new_callback])

    # Connect the model and restore callbacks before evaluation
    new_trainer.strategy.connect(new_model)
    new_trainer._checkpoint_connector.resume_start(ckpt_path)
    new_trainer._checkpoint_connector.restore_callbacks()

    # Run the evaluation phase (validate/test/predict)
    fn = getattr(new_trainer, trainer_fn)
    fn(new_model, ckpt_path=ckpt_path)

    assert new_callback.restored  # Should be True after loading the checkpoint
Test that callbacks are properly restored in non-fit phases.
tests/tests_pytorch/trainer/connectors/test_checkpoint_connector.py