def test_checkpoint_plugin_called(tmp_path):
    """Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when saving/loading."""
    checkpoint_plugin = CustomCheckpointIO()
    checkpoint_plugin = MagicMock(wraps=checkpoint_plugin, spec=CustomCheckpointIO)

    ck = ModelCheckpoint(dirpath=tmp_path, save_last=True)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cpu",
        strategy=SingleDeviceStrategy("cpu", checkpoint_io=checkpoint_plugin),
        callbacks=ck,
        max_epochs=2,
        limit_train_batches=1,
        limit_val_batches=0,
        limit_test_batches=1,
    )
    trainer.fit(model)

    ckpt_files = {fn.name for fn in Path(tmp_path).glob("*.ckpt")}
    assert ckpt_files == {"epoch=1-step=2.ckpt", "last.ckpt"}
    assert trainer.checkpoint_callback.best_model_path == str(tmp_path / "epoch=1-step=2.ckpt")
    assert trainer.checkpoint_callback.last_model_path == str(tmp_path / "last.ckpt")
    assert checkpoint_plugin.save_checkpoint.call_count == 4
    assert checkpoint_plugin.remove_checkpoint.call_count == 1

    trainer.test(model, ckpt_path=ck.last_model_path)
    checkpoint_plugin.load_checkpoint.assert_called_with(str(tmp_path / "last.ckpt"))

    checkpoint_plugin.reset_mock()
    ck = ModelCheckpoint(dirpath=tmp_path, save_last=True)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cpu",
        strategy=SingleDeviceStrategy("cpu"),
        plugins=[checkpoint_plugin],
        callbacks=ck,
        max_epochs=2,
        limit_train_batches=1,
        limit_val_batches=0,
        limit_test_batches=1,
    )
    trainer.fit(model)

    ckpt_files = {fn.name for fn in Path(tmp_path).glob("*.ckpt")}
    assert ckpt_files == {"epoch=1-step=2.ckpt", "last.ckpt", "epoch=1-step=2-v1.ckpt", "last-v1.ckpt"}
    assert trainer.checkpoint_callback.best_model_path == str(tmp_path / "epoch=1-step=2-v1.ckpt")
    assert trainer.checkpoint_callback.last_model_path == str(tmp_path / "last-v1.ckpt")
    assert checkpoint_plugin.save_checkpoint.call_count == 4
    assert checkpoint_plugin.remove_checkpoint.call_count == 1

    trainer.test(model, ckpt_path=ck.last_model_path)
    checkpoint_plugin.load_checkpoint.assert_called_once()
    checkpoint_plugin.load_checkpoint.assert_called_with(str(tmp_path / "last-v1.ckpt"))
Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when saving/loading.
tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py