code: def test_lightning_dataset(get_dataset, strategy_type):
      import pytorch_lightning as pl
  
      dataset = get_dataset(name="MUTAG").shuffle()
      train_dataset = dataset[:50]
      val_dataset = dataset[50:80]
      test_dataset = dataset[80:90]
      pred_dataset = dataset[90:]
  
      devices = (
          1 if strategy_type is None else torch.cuda.device_count()
      )
      if strategy_type == "ddp":
          strategy = pl.strategies.DDPStrategy(accelerator="gpu")
      else:
          strategy = pl.strategies.SingleDeviceStrategy(device="cuda:0")
  
      model = LinearGraphModule(
          dataset.num_features, 64, dataset.num_classes
      )
  
      trainer = pl.Trainer(
          strategy=strategy,
          devices=devices,
          max_epochs=1,
          log_every_n_steps=1,
      )
      with pytest.warns(
          UserWarning, match="'shuffle=True' option is ignored"
      ):
          datamodule = LightningDataset(
              train_dataset,
              val_dataset,
              test_dataset,
              pred_dataset,
              batch_size=5,
              num_workers=3,
              shuffle=True,
          )
          assert "shuffle" not in datamodule.kwargs
      old_x = train_dataset._data.x.clone()
      assert str(datamodule) == (
          "LightningDataset(train_dataset=MUTAG(50), "
          "val_dataset=MUTAG(30), "
          "test_dataset=MUTAG(10), "
          "pred_dataset=MUTAG(98), batch_size=5, "
          "num_workers=3, pin_memory=True, "
          "persistent_workers=True)"
      )
      trainer.fit(model, datamodule)
      trainer.test(model, datamodule)
      new_x = train_dataset._data.x
      assert torch.all(new_x > old_x)  # Ensure shared data.
      assert trainer.validate_loop._data_source.is_defined()
      assert trainer.test_loop._data_source.is_defined()
  
      # Test with `val_dataset=None` and `test_dataset=None`:
      if strategy_type is None:
          trainer = pl.Trainer(
              strategy=strategy,
              devices=devices,
              max_epochs=1,
              log_every_n_steps=1,
          )
  
          datamodule = LightningDataset(train_dataset, batch_size=5)
          assert str(datamodule) == (
              "LightningDataset(train_dataset=MUTAG(50), "
              "batch_size=5, num_workers=0, "
              "pin_memory=True, "
              "persistent_workers=False)"
          )
  
          with pytest.warns(
              UserWarning, match="defined a `validation_step`"
          ):
              trainer.fit(model, datamodule)
  
          assert not trainer.validate_loop._data_source.is_defined()
          assert not trainer.test_loop._data_source.is_defined()
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.data.lightning import LightningDataset, LightningLinkData, LightningNodeData
  from torch_geometric.nn import global_mean_pool
  from torch_geometric.sampler import BaseSampler, NeighborSampler
  from torch_geometric.testing import MyFeatureStore, MyGraphStore, get_random_edge_index, onlyCUDA, onlyFullTest, onlyNeighborSampler, onlyOnline, withPackage
label_desc: Use image node_id 1 for calling a global function with example usage: test_lightning_dataset(get_dataset, strategy_type) without return types