Fix tests?

Author: Dsantra92

test/datasets/graphs.jl

@@ -11,9 +11,9 @@
     @test g.num_edges == 9104
     @test size(g.node_data.features) == (3703, g.num_nodes)
     @test size(g.node_data.targets) == (g.num_nodes,)
-    @test sum(g.node_data.train_mask) == 120
-    @test sum(g.node_data.val_mask) == 500
-    @test sum(g.node_data.test_mask) == 1015
+    @test sum(data.metadata["node"].split[1].train) == 120
+    @test sum(data.metadata["node"].split[1].val) == 500
+    @test sum(data.metadata["node"].split[1].test) == 1015
     @test g.edge_index isa Tuple{Vector{Int}, Vector{Int}}
     s, t = g.edge_index
     for a in (s, t)
@@ -36,9 +36,9 @@ end
     @test g.num_edges == 10556
     @test size(g.node_data.features) == (1433, g.num_nodes)
     @test size(g.node_data.targets) == (g.num_nodes,)
-    @test sum(g.node_data.train_mask) == 140
-    @test sum(g.node_data.val_mask) == 500
-    @test sum(g.node_data.test_mask) == 1000
+    @test sum(data.metadata["node"].split[1].train) == 140
+    @test sum(data.metadata["node"].split[1].val) == 500
+    @test sum(data.metadata["node"].split[1].test) == 1000
     @test g.edge_index isa Tuple{Vector{Int}, Vector{Int}}
     s, t = g.edge_index
     for a in (s, t)
@@ -109,9 +109,9 @@ end
     @test g.num_edges == 88648
     @test size(g.node_data.features) == (500, g.num_nodes)
     @test size(g.node_data.targets) == (g.num_nodes,)
-    @test sum(g.node_data.train_mask) == 60
-    @test sum(g.node_data.val_mask) == 500
-    @test sum(g.node_data.test_mask) == 1000
+    @test sum(data.metadata["node"].split[1].train) == 60
+    @test sum(data.metadata["node"].split[1].val) == 500
+    @test sum(data.metadata["node"].split[1].test) == 1000
     @test g.edge_index isa Tuple{Vector{Int}, Vector{Int}}
     s, t = g.edge_index
     for a in (s, t)

test/datasets/graphs_no_ci.jl

@@ -247,13 +247,21 @@ end
     @test g.num_nodes == 169343
     @test g.num_edges == 1166243
 
-    @test sum(count.([g.node_data.train_mask, g.node_data.test_mask, g.node_data.val_mask])) == g.num_nodes
+    train_mask = d.metadata["node"].split["time"][1].train
+    test_mask = d.metadata["node"].split["time"][1].test
+    val_mask = d.metadata["node"].split["time"][1].val
+
+    @test sum(count.([train_mask, test_mask, val_mask])) == g.num_nodes
 end
 
 @testset "OGBDataset - ogbg-molhiv" begin
     d = OGBDataset("ogbg-molhiv")
 
-    @test sum(count.([d.graph_data.train_mask, d.graph_data.test_mask, d.graph_data.val_mask])) == length(d)
+    train_mask = d.metadata["graph"].split["scaffold"].train
+    test_mask = d.metadata["graph"].split["scaffold"].test
+    val_mask = d.metadata["graph"].split["scaffold"].val
+
+    @test sum(count.([train_mask, test_mask, val_mask])) == length(d)
 end
 
 @testset "Reddit_full" begin
@@ -264,9 +272,10 @@ end
     @test g.num_edges == 114615892
     @test size(g.node_data.features) == (602, g.num_nodes)
     @test size(g.node_data.labels) == (g.num_nodes,)
-    @test count(g.node_data.train_mask) == 153431
-    @test count(g.node_data.val_mask) == 23831
-    @test count(g.node_data.test_mask) == 55703
+    split = data.metadata["node"].split[1]
+    @test count(split.train) == 153431
+    @test count(split.val) == 23831
+    @test count(split.test) == 55703
     s, t = g.edge_index
     @test length(s) == length(t) == g.num_edges
     @test minimum(s) == minimum(t) == 1
@@ -281,9 +290,10 @@ end
     @test g.num_edges == 23213838
     @test size(g.node_data.features) == (602, g.num_nodes)
     @test size(g.node_data.labels) == (g.num_nodes,)
-    @test count(g.node_data.train_mask) == 152410
-    @test count(g.node_data.val_mask) == 23699
-    @test count(g.node_data.test_mask) == 55334
+    split = data.metadata["node"].split[1]
+    @test count(split.train) == 152410
+    @test count(split.val) == 23699
+    @test count(split.test) == 55334
     s, t = g.edge_index
     @test length(s) == length(t) == g.num_edges
     @test minimum(s) == minimum(t) == 1