code: def test_init():
      train_input = DatasetInput(RunningStage.TRAINING, DummyDataset())
      val_input = DatasetInput(RunningStage.VALIDATING, DummyDataset())
      test_input = DatasetInput(RunningStage.TESTING, DummyDataset())
  
      data_module = DataModule(train_input, batch_size=1)
      assert data_module.train_dataset
      assert not data_module.val_dataset
      assert not data_module.test_dataset
  
      data_module = DataModule(train_input, val_input, batch_size=1)
      assert data_module.train_dataset
      assert data_module.val_dataset
      assert not data_module.test_dataset
  
      data_module = DataModule(
          train_input, val_input, test_input, batch_size=1
      )
      assert data_module.train_dataset
      assert data_module.val_dataset
      assert data_module.test_dataset
  
file_path: lightning-flash/tests/core/test_data.py
imports: import pytest
  import torch
  from flash import DataKeys, DataModule, RunningStage
  from flash.core.data.data_module import DatasetInput
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE
label_desc: Use image node_id 1 for calling a global function with example usage: test_init() without return types