From 085ab6d65da02a9e456ba9da31228b95b99834f8 Mon Sep 17 00:00:00 2001 From: kyrylo Date: Thu, 8 Dec 2022 11:33:01 +0100 Subject: [PATCH] Made requested changes --- tests/web/test_optimizer_api.py | 2 +- txmatching/optimizer/optimizer_functions.py | 14 +++++++------- txmatching/scorers/compatibility_graph.py | 5 ++++- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/web/test_optimizer_api.py b/tests/web/test_optimizer_api.py index 07eb99717..560932f8d 100644 --- 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 index ddbf7e139..9b5f3969e 100644 --- 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 index 8f61fd273..0b4f4bb56 100644 --- 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]