From a72b6b09bcf5765847dd956e1d1b4cf7486fd585 Mon Sep 17 00:00:00 2001
From: Stanislaw Adaszewski <stanislaw.adaszewski@roche.com>
Date: Thu, 11 Jun 2020 18:37:32 +0200
Subject: [PATCH] Fix regressions in test_input.

---
 tests/icosagon/test_input.py | 29 +++++++++--------------------
 1 file changed, 9 insertions(+), 20 deletions(-)

diff --git a/tests/icosagon/test_input.py b/tests/icosagon/test_input.py
index 1e9cfbf..1ac2e3d 100644
--- a/tests/icosagon/test_input.py
+++ b/tests/icosagon/test_input.py
@@ -9,28 +9,17 @@ def _some_data():
     d = Data()
     d.add_node_type('Gene', 1000)
     d.add_node_type('Drug', 100)
-    d.add_relation_type('Target', 1, 0, torch.rand(100, 1000))
-    d.add_relation_type('Interaction', 0, 0, torch.rand(1000, 1000))
-    d.add_relation_type('Side Effect: Nausea', 1, 1, torch.rand(100, 100))
-    d.add_relation_type('Side Effect: Infertility', 1, 1, torch.rand(100, 100))
-    d.add_relation_type('Side Effect: Death', 1, 1, torch.rand(100, 100))
-    return d
 
+    fam = d.add_relation_family('Drug-Gene', 1, 0, False)
+    fam.add_relation_type('Target', 1, 0, torch.rand(100, 1000))
 
-def _some_data_with_interactions():
-    d = Data()
-    d.add_node_type('Gene', 1000)
-    d.add_node_type('Drug', 100)
-    d.add_relation_type('Target', 1, 0,
-        torch.rand((100, 1000), dtype=torch.float32).round())
-    d.add_relation_type('Interaction', 0, 0,
-        torch.rand((1000, 1000), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Nausea', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Infertility', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Death', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
+    fam = d.add_relation_family('Gene-Gene', 0, 0, False)
+    fam.add_relation_type('Interaction', 0, 0, torch.rand(1000, 1000))
+
+    fam = d.add_relation_family('Drug-Drug', 1, 1, False)
+    fam.add_relation_type('Side Effect: Nausea', 1, 1, torch.rand(100, 100))
+    fam.add_relation_type('Side Effect: Infertility', 1, 1, torch.rand(100, 100))
+    fam.add_relation_type('Side Effect: Death', 1, 1, torch.rand(100, 100))
     return d