rtest for repr

pyg-team · Sep 15, 2022 · ef98779 · ef98779
1 parent a054f15
commit ef98779
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/test/data/test_summary.py b/test/data/test_summary.py
@@ -27,6 +27,7 @@ def test_summary():
     seed_everything(0)
     dataset = FakeDataset(num_graphs=10)
     summary = dataset.summary()
+    assert "FakeDataset" in str(summary)
     assert summary.num_graphs == 10
     num_nodes = torch.tensor([d.num_nodes for d in dataset]).float()
     num_edges = torch.tensor([d.num_edges for d in dataset]).float()
@@ -50,6 +51,7 @@ def test_hetero_summary():
     seed_everything(0)
     dataset = FakeHeteroDataset(num_graphs=10)
     summary = dataset.summary()
+    assert "FakeHeteroDataset" in str(summary)
     assert summary.num_graphs == 10
     num_nodes = torch.tensor([d.num_nodes for d in dataset]).float()
     num_edges = torch.tensor([d.num_edges for d in dataset]).float()