code: def test_data_module():
      seed_everything(42)
  
      def train_fn(data):
          return data - 100
  
      def val_fn(data):
          return data + 100
  
      def test_fn(data):
          return data - 1000
  
      def predict_fn(data):
          return data + 1000
  
      @dataclass
      class TestTransform(InputTransform):
          def per_sample_transform(self):
              def fn(x):
                  return x
  
              return fn
  
          def train_per_batch_transform_on_device(self) -> Callable:
              return train_fn
  
          def val_per_batch_transform_on_device(self) -> Callable:
              return val_fn
  
          def test_per_batch_transform_on_device(self) -> Callable:
              return test_fn
  
          def predict_per_batch_transform_on_device(self) -> Callable:
              return predict_fn
  
      transform = TestTransform()
      assert transform._transform is not None
  
      train_dataset = Input(
          RunningStage.TRAINING, np.arange(10, dtype=np.float32)
      )
      assert train_dataset.running_stage == RunningStage.TRAINING
  
      val_dataset = Input(
          RunningStage.VALIDATING, np.arange(10, dtype=np.float32)
      )
      assert val_dataset.running_stage == RunningStage.VALIDATING
  
      test_dataset = Input(
          RunningStage.TESTING, np.arange(10, dtype=np.float32)
      )
      assert test_dataset.running_stage == RunningStage.TESTING
  
      predict_dataset = Input(
          RunningStage.PREDICTING, np.arange(10, dtype=np.float32)
      )
      assert predict_dataset.running_stage == RunningStage.PREDICTING
  
      dm = DataModule(
          train_input=train_dataset,
          val_input=val_dataset,
          test_input=test_dataset,
          predict_input=predict_dataset,
          transform=transform,
          batch_size=2,
      )
      assert len(dm.train_dataloader()) == 5
      batch = next(iter(dm.train_dataloader()))
      assert batch.shape == torch.Size([2])
      assert batch.min() >= 0
      assert batch.max() < 10
  
      assert len(dm.val_dataloader()) == 5
      batch = next(iter(dm.val_dataloader()))
      assert batch.shape == torch.Size([2])
      assert batch.min() >= 0
      assert batch.max() < 10
  
      class TestModel(Task):
          def training_step(self, batch, batch_idx):
              assert sum(batch < 0) == 2
  
          def validation_step(self, batch, batch_idx):
              assert sum(batch > 0) == 2
  
          def test_step(self, batch, batch_idx):
              assert sum(batch < 500) == 2
  
          def predict_step(self, batch, *args, **kwargs):
              assert sum(batch > 500) == 2
              assert torch.equal(batch, torch.tensor([1000.0, 1001.0]))
  
          def on_train_dataloader(self) -> None:
              pass
  
          def on_val_dataloader(self) -> None:
              pass
  
          def on_test_dataloader(self, *_) -> None:
              pass
  
          def on_predict_dataloader(self) -> None:
              pass
  
          def on_predict_end(self) -> None:
              pass
  
          def on_fit_end(self) -> None:
              pass
  
      model = TestModel(torch.nn.Linear(1, 1))
      trainer = Trainer(fast_dev_run=True)
      trainer.fit(model, datamodule=dm)
      trainer.validate(model, datamodule=dm)
      trainer.test(model, datamodule=dm)
      trainer.predict(model, datamodule=dm)
  
      # Test that plain lightning module works with FlashDataModule
      class SampleBoringModel(BoringModel):
          def __init__(self):
              super().__init__()
              self.layer = torch.nn.Linear(2, 1)
  
      model = SampleBoringModel()
      trainer = Trainer(fast_dev_run=True)
      trainer.fit(model, datamodule=dm)
      trainer.validate(model, datamodule=dm)
      trainer.test(model, datamodule=dm)
      trainer.predict(model, datamodule=dm)
  
      transform = TestTransform()
      input = Input(RunningStage.TRAINING)
      dm = DataModule(
          train_input=input, batch_size=1, transform=transform
      )
      assert isinstance(dm.input_transform, TestTransform)
  
      class RandomDataset(Dataset):
          def __init__(self, size: int, length: int):
              self.len = length
              self.data = torch.ones(length, size)
  
          def __getitem__(self, index):
              return self.data[index]
  
          def __len__(self):
              return self.len
  
      def _add_hundred(x):
          if isinstance(x, Dict):
              x["input"] += 100
          else:
              x += 100
          return x
  
      class TrainInputTransform(InputTransform):
          def _add_one(self, x):
              if isinstance(x, Dict):
                  x["input"] += 1
              else:
                  x += 1
              return x
  
          def per_sample_transform(self) -> Callable:
              return self._add_one
  
          def val_per_sample_transform(self) -> Callable:
              return _add_hundred
  
      dm = DataModule(
          train_input=DatasetInput(
              RunningStage.TRAINING, RandomDataset(64, 32)
          ),
          val_input=DatasetInput(
              RunningStage.VALIDATING, RandomDataset(64, 32)
          ),
          test_input=DatasetInput(
              RunningStage.TESTING, RandomDataset(64, 32)
          ),
          batch_size=3,
          transform=TrainInputTransform(),
      )
      batch = next(iter(dm.train_dataloader()))
      assert batch["input"][0][0] == 2
      batch = next(iter(dm.val_dataloader()))
      assert batch["input"][0][0] == 101
      batch = next(iter(dm.test_dataloader()))
      assert batch["input"][0][0] == 2
  
imports: from dataclasses import dataclass
  from typing import Callable, Dict
  from unittest.mock import MagicMock, NonCallableMock, patch
  import numpy
  import pytest
  import torch
  from flash import Task, Trainer
  from flash.core.data.data_module import DataModule, DatasetInput
  from flash.core.data.io.input import Input
  from flash.core.data.io.input_transform import InputTransform
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE, _TOPIC_IMAGE_AVAILABLE, _TORCHVISION_AVAILABLE
  from flash.core.utilities.stages import RunningStage
  from pytorch_lightning import seed_everything
  from torch.utils.data import Dataset
  from tests.helpers.boring_model import BoringModel
label_desc: Use image node_id 1 for calling a global function with example usage: test_data_module() and returns: unknown, unknown, unknown, unknown, x, fn, train_fn, val_fn, test_fn, predict_fn, self, self, x, self, _add_hundred, x