Made requested changes

mild-blue · Dec 13, 2022 · 085ab6d · 085ab6d
1 parent 2fc4fd8
commit 085ab6d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 9 deletions.
diff --git a/tests/web/test_optimizer_api.py b/tests/web/test_optimizer_api.py
@@ -146,7 +146,7 @@ def test_optimizer_export_api_works(self):
 
         compatibility_graph = get_compatibility_graph_for_optimizer_api(txm_event.active_and_valid_donors_dict,
                                                                         txm_event.active_and_valid_recipients_dict)
-        self.assertCountEqual(compatibility_graph, res.json['compatibility_graph'])
+        self.assertEqual(len(compatibility_graph), len(res.json['compatibility_graph']))
 
         # we are calling txm event with default configuration
         def_config = ConfigParameters()

diff --git a/txmatching/optimizer/optimizer_functions.py b/txmatching/optimizer/optimizer_functions.py
@@ -262,7 +262,7 @@ def get_optimizer_configuration(config: ConfigParameters) -> OptimizerConfigurat
 
 def get_compatibility_graph_for_optimizer_api(donors_dict: Dict[DonorDbId, Donor],
                                               recipients_dict: Dict[
-                                                  RecipientDbId, Recipient]) -> OptimizerCompatibilityGraph:
+                                              RecipientDbId, Recipient]) -> OptimizerCompatibilityGraph:
     scorer = SplitScorer()
 
     compatibility_graph_from_scorer = scorer.get_compatibility_graph(recipients_dict, donors_dict)
@@ -271,12 +271,12 @@ def get_compatibility_graph_for_optimizer_api(donors_dict: Dict[DonorDbId, Donor
     for i, donor_id in enumerate(donors_dict):
         for j, recipient_id in enumerate(recipients_dict):
             if (i, j) in compatibility_graph_from_scorer:
-                comp_graph_cell = {
-                    "donor_id": donor_id,
-                    "recipient_id": recipient_id,
-                    "weights": {
+                comp_graph_entry = CompatibilityGraphEntry(
+                    donor_id=donor_id,
+                    recipient_id=recipient_id,
+                    weights={
                         "hla_compatibility_score": int(compatibility_graph_from_scorer[(i, j)]['score'])
                     }
-                }
-                compatibility_graph.append(comp_graph_cell)
+                )
+                compatibility_graph.append(comp_graph_entry)
     return compatibility_graph
diff --git a/txmatching/scorers/compatibility_graph.py b/txmatching/scorers/compatibility_graph.py
@@ -1,5 +1,8 @@
 from typing import Dict, List, Tuple
 
+from txmatching.optimizer.optimizer_request_object import \
+    CompatibilityGraphEntry
+
 CompatibilityGraph = Dict[Tuple[int, int], Dict[str, int]]
 
-OptimizerCompatibilityGraph = List[Dict[str, any]]
+OptimizerCompatibilityGraph = List[CompatibilityGraphEntry]