From 5c24883ba4eaf4df27bd03ea72e547d7b1ea6c76 Mon Sep 17 00:00:00 2001
From: Stanislaw Adaszewski <stanislaw.adaszewski@roche.com>
Date: Sat, 20 Jun 2020 17:56:20 +0200
Subject: [PATCH] Add test_prep_rel_two_node_types_sym_01().

---
 tests/icosagon/test_trainprep.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/tests/icosagon/test_trainprep.py b/tests/icosagon/test_trainprep.py
index 2efcf45..79a8aee 100644
--- a/tests/icosagon/test_trainprep.py
+++ b/tests/icosagon/test_trainprep.py
@@ -9,7 +9,8 @@ from icosagon.trainprep import TrainValTest, \
     get_edges_and_degrees, \
     prepare_adj_mat, \
     prepare_relation_type, \
-    prep_rel_one_node_type
+    prep_rel_one_node_type, \
+    prep_rel_two_node_types_sym
 import torch
 import pytest
 import numpy as np
@@ -139,6 +140,33 @@ def test_prep_rel_one_node_type_01():
     assert len(prep_rel.edges_back_neg.test) == 0
 
 
+def test_prep_rel_two_node_types_sym_01():
+    adj_mat = torch.zeros(200)
+    perm = torch.randperm(100)
+    adj_mat[perm[:10]] = 1
+    adj_mat = adj_mat.view(10, 20)
+    rel = RelationType('Dummy Relation', 0, 1, adj_mat, None)
+    ratios = TrainValTest(.8, .1, .1)
+    prep_rel = prep_rel_two_node_types_sym(rel, ratios)
+    assert prep_rel.name == rel.name
+    assert prep_rel.node_type_row == rel.node_type_row
+    assert prep_rel.node_type_column == rel.node_type_column
+    assert prep_rel.adjacency_matrix.shape == rel.adjacency_matrix.shape
+    assert prep_rel.adjacency_matrix_backward.shape == (20, 10)
+    assert len(prep_rel.edges_pos.train) == 8
+    assert len(prep_rel.edges_pos.val) == 1
+    assert len(prep_rel.edges_pos.test) == 1
+    assert len(prep_rel.edges_neg.train) == 8
+    assert len(prep_rel.edges_neg.val) == 1
+    assert len(prep_rel.edges_neg.test) == 1
+
+    assert len(prep_rel.edges_back_pos.train) == 0
+    assert len(prep_rel.edges_back_pos.val) == 0
+    assert len(prep_rel.edges_back_pos.test) == 0
+    assert len(prep_rel.edges_back_neg.train) == 0
+    assert len(prep_rel.edges_back_neg.val) == 0
+    assert len(prep_rel.edges_back_neg.test) == 0
+
 
 # def prepare_relation(r, ratios):
 #     adj_mat = r.adjacency_matrix