def test_flash_callback(_, __, tmpdir):  # noqa: PT019
    """Test the callback hook system for fit."""

    callback_mock = MagicMock()

    inputs = [(torch.rand(1), torch.rand(1))]
    transform = InputTransform()
    dm = DataModule(
        DatasetInput(RunningStage.TRAINING, inputs),
        DatasetInput(RunningStage.VALIDATING, inputs),
        DatasetInput(RunningStage.TESTING, inputs),
        transform=transform,
        batch_size=1,
        num_workers=0,
        data_fetcher=callback_mock,
    )

    _ = next(iter(dm.train_dataloader()))

    assert callback_mock.method_calls == [
        call.on_load_sample(ANY, RunningStage.TRAINING),
        call.on_per_sample_transform(ANY, RunningStage.TRAINING),
        call.on_collate(ANY, RunningStage.TRAINING),
        call.on_per_batch_transform(ANY, RunningStage.TRAINING),
    ]

    class CustomModel(Task):
        def __init__(self):
            super().__init__(model=torch.nn.Linear(1, 1), loss_fn=torch.nn.MSELoss())

        def training_step(self, batch, batch_idx):
            batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
            return super().training_step(batch, batch_idx)

        def validation_step(self, batch, batch_idx):
            batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
            return super().validation_step(batch, batch_idx)

        def test_step(self, batch, batch_idx):
            batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
            return super().test_step(batch, batch_idx)

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        limit_val_batches=1,
        limit_train_batches=1,
    )
    transform = InputTransform()
    dm = DataModule(
        DatasetInput(RunningStage.TRAINING, inputs),
        DatasetInput(RunningStage.VALIDATING, inputs),
        DatasetInput(RunningStage.TESTING, inputs),
        transform=transform,
        batch_size=1,
        num_workers=0,
        data_fetcher=callback_mock,
    )
    trainer.fit(CustomModel(), datamodule=dm)

    assert callback_mock.method_calls == [
        call.on_load_sample(ANY, RunningStage.TRAINING),
        call.on_per_sample_transform(ANY, RunningStage.TRAINING),
        call.on_collate(ANY, RunningStage.TRAINING),
        call.on_per_batch_transform(ANY, RunningStage.TRAINING),
        call.on_load_sample(ANY, RunningStage.VALIDATING),
        call.on_per_sample_transform(ANY, RunningStage.VALIDATING),
        call.on_collate(ANY, RunningStage.VALIDATING),
        call.on_per_batch_transform(ANY, RunningStage.VALIDATING),
        call.on_per_batch_transform_on_device(ANY, RunningStage.VALIDATING),
        call.on_load_sample(ANY, RunningStage.TRAINING),
        call.on_per_sample_transform(ANY, RunningStage.TRAINING),
        call.on_collate(ANY, RunningStage.TRAINING),
        call.on_per_batch_transform(ANY, RunningStage.TRAINING),
        call.on_per_batch_transform_on_device(ANY, RunningStage.TRAINING),
        call.on_load_sample(ANY, RunningStage.VALIDATING),
        call.on_per_sample_transform(ANY, RunningStage.VALIDATING),
        call.on_collate(ANY, RunningStage.VALIDATING),
        call.on_per_batch_transform(ANY, RunningStage.VALIDATING),
        call.on_per_batch_transform_on_device(ANY, RunningStage.VALIDATING),
    ]
Test the callback hook system for fit.
tests/core/data/test_callback.py