def test_request_dataloader(tmp_path):
    """This test asserts dataloader can be wrapped."""

    class DataLoaderWrapper:
        def __init__(self, loader):
            self.loader = loader
            self._iter = iter(self.loader)

        def __iter__(self):
            self._iter = iter(self.loader)
            return self._iter

        def __next__(self):
            return next(self._iter)

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.on_train_batch_start_called = False
            self.on_val_batch_start_called = False

        def train_dataloader(self):
            loader = super().train_dataloader()
            return DataLoaderWrapper(loader)

        def on_train_batch_start(self, batch, batch_idx: int) -> None:
            assert isinstance(self.trainer.train_dataloader, DataLoaderWrapper)
            self.on_train_batch_start_called = True

        def val_dataloader(self):
            loader = super().val_dataloader()
            return DataLoaderWrapper(loader)

        def on_validation_batch_start(self, *_):
            assert isinstance(self.trainer.val_dataloaders, DataLoaderWrapper)
            self.on_val_batch_start_called = True

    trainer = Trainer(
        default_root_dir=tmp_path, limit_train_batches=2, limit_val_batches=2, limit_test_batches=2, max_epochs=1
    )
    model = TestModel()
    trainer.fit(model)
    trainer.test(model)
    assert model.on_train_batch_start_called
    assert model.on_val_batch_start_called
This test asserts dataloader can be wrapped.
tests/tests_pytorch/trainer/test_dataloaders.py