Simplify test

Author: aurorarossi

test/utils.jl

@@ -61,41 +61,31 @@
     end
 
     @testset "topk_nodes" begin
-        A = [0.0297 0.8307 0.9140 0.6702 0.3346;
-             0.5901 0.3030 0.9280 0.6893 0.7997;
-             0.0880 0.6515 0.4451 0.7507 0.5297;
-             0.5171 0.6379 0.2695 0.8954 0.5197]
-        B = [0.3168 0.3174 0.5303 0.0804 0.3808;
-             0.1752 0.9105 0.5692 0.8489 0.0539;
-             0.1931 0.4954 0.3455 0.3934 0.0857;
-             0.5065 0.5182 0.5418 0.1520 0.3872]
-        C = [0.0297 0.0297 0.8307 0.9140 0.6702 0.3346;
-             0.5901 0.5901 0.3030 0.9280 0.6893 0.7997;
-             0.0880 0.0880 0.6515 0.4451 0.7507 0.5297;
-             0.5171 0.5171 0.6379 0.2695 0.8954 0.5197]
-        g1 = rand_graph(5, 6, ndata = (w = A,))
-        g2 = rand_graph(5, 6, ndata = (w = B,))
-        g3 = rand_graph(5, 6, edata = (e = C,))
+        A = [1.0 5.0 9.0; 2.0 6.0 10.0; 3.0 7.0 11.0; 4.0 8.0 12.0]
+        B = [0.318907 0.189981 0.991791;
+             0.547022 0.977349 0.680538;
+             0.921823 0.35132 0.494715;
+             0.451793 0.00704976 0.0189275]
+        g1 = rand_graph(3, 6, ndata = (x = A,))
+        g2 = rand_graph(3, 6, ndata = B)
+
+        #
+        output1 = topk_nodes(g1, :x, 2)
+        output2 = topk_nodes(g2, :x, 1, sortby = 2)
+
+        @test output1 == [9.0 5.0;
+                          10.0 6.0;
+                          11.0 7.0;
+                          12.0 8.0]
+        @test output2 == [0.189981;
+                          0.977349;
+                          0.35132;
+                          0.00704976;;]
         g = Flux.batch([g1, g2])
-        output1 = topk_nodes(g1, :w, 3)
-        output2 = topk_nodes(g1, :w, 3; sortby = 5)
-        output3 = topk_edges(g3, :e, 3; sortby = 6)
-        output_batch = topk_nodes(g, :w, 3; sortby = 5)
-        correctout1 = [0.5901 0.8307 0.9280 0.8954 0.7997;
-                       0.5171 0.6515 0.9140 0.7507 0.5297;
-                       0.0880 0.6379 0.4451 0.6893 0.5197]
-        correctout2 = [0.5901 0.3030 0.9280 0.6893 0.7997;
-                       0.0880 0.6515 0.4451 0.7507 0.5297;
-                       0.5171 0.6379 0.2695 0.8954 0.5197]
-        correctout3 = [0.5901 0.5901 0.3030 0.9280 0.6893 0.7997;
-                       0.0880 0.0880 0.6515 0.4451 0.7507 0.5297;
-                       0.5171 0.5171 0.6379 0.2695 0.8954 0.5197]
-        correctout_batch = [0.5901 0.3030 0.9280 0.6893 0.7997 0.5065 0.5182 0.5418 0.1520 0.3872;
-                            0.0880 0.6515 0.4451 0.7507 0.5297 0.3168 0.3174 0.5303 0.0804 0.3808;
-                            0.5171 0.6379 0.2695 0.8954 0.5197 0.1931 0.4954 0.3455 0.3934 0.0857]
-        @test output1 == correctout1
-        @test output2 == correctout2
-        @test output3 == correctout3
-        @test output_batch == correctout_batch
+        output3 = topk_nodes(g, :x, 2; sortby = 4)
+        @test output3 == [9.0 5.0 0.318907 0.991791;
+               10.0 6.0 0.547022 0.680538;
+               11.0 7.0 0.921823 0.494715;
+               12.0 8.0 0.451793 0.0189275]
     end
-end
+end;