def test_async_checkpoint_plugin(tmp_path):
    """Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when async saving and
    loading."""

    checkpoint_plugin = AsyncCheckpointIO()

    checkpoint_plugin.save_checkpoint = Mock(wraps=checkpoint_plugin.save_checkpoint)
    checkpoint_plugin.remove_checkpoint = Mock(wraps=checkpoint_plugin.remove_checkpoint)

    class CustomBoringModel(BoringModel):
        def on_fit_start(self):
            base_ckpt_io = self.trainer.strategy.checkpoint_io.checkpoint_io
            base_ckpt_io.save_checkpoint = Mock(wraps=base_ckpt_io.save_checkpoint)
            base_ckpt_io.remove_checkpoint = Mock(wraps=base_ckpt_io.remove_checkpoint)

    ck = ModelCheckpoint(dirpath=tmp_path, save_top_k=2, monitor="step", mode="max")

    model = CustomBoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        plugins=[checkpoint_plugin],
        callbacks=ck,
        max_epochs=3,
        limit_train_batches=1,
        limit_val_batches=0,
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.fit(model)

    assert checkpoint_plugin.save_checkpoint.call_count == 3
    assert checkpoint_plugin.remove_checkpoint.call_count == 1

    base_ckpt_io = trainer.strategy.checkpoint_io.checkpoint_io
    assert base_ckpt_io.save_checkpoint.call_count == 3
    assert base_ckpt_io.remove_checkpoint.call_count == 1
Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when async saving and
    loading.
tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py