code: def test_infection_dataset():
      seed_everything(12345)
      graph_generator = DummyGraph()
      dataset = InfectionDataset(
          graph_generator, num_infected_nodes=2, max_path_length=2
      )
      assert str(dataset) == (
          "InfectionDataset(1, "
          "graph_generator=DummyGraph(), "
          "num_infected_nodes=2, "
          "max_path_length=2)"
      )
      assert len(dataset) == 1
  
      data = dataset[0]
      assert len(data) == 4
      assert data.x.size() == (10, 2)
      assert data.x[:, 0].sum() == 8 and data.x[:, 1].sum() == 2
      assert torch.equal(data.edge_index, graph_generator().edge_index)
      assert data.y.size() == (10,)
  
      # With `seed=12345`, node 0 and node 7 will be infected:
      assert data.x[0].tolist() == [0, 1]  # First infected node.
      assert data.x[7].tolist() == [0, 1]  # Second infected node.
      assert data.y.tolist() == [0, 1, 2, 3, 3, 2, 1, 0, 1, 2]
      assert data.edge_mask.tolist() == [
          1,
          1,
          1,
          0,
          0,
          0,
          0,
          0,
          0,
          0,
          0,
          1,
          1,
          1,
          1,
          1,
          1,
          0,
      ]
  
file_path: pytorch_geometric/test/datasets/test_infection_dataset.py
imports: import torch
  from torch_geometric import seed_everything
  from torch_geometric.data import Data
  from torch_geometric.datasets import InfectionDataset
  from torch_geometric.datasets.graph_generator import ERGraph, GraphGenerator
label_desc: Use image node_id 1 for calling a global function with example usage: test_infection_dataset() without return types