From 6f4fd59c945d1a53a0be1c373c2c37f6dd1d8c7d Mon Sep 17 00:00:00 2001
From: Alexander Berger <alexander.berger@jax.org>
Date: Tue, 3 Sep 2024 14:56:45 -0400
Subject: [PATCH] Fixing tests and doing some error checking

---
 src/geneweaver/client/api/mapping.py |  5 ++-
 tests/unit/api/test_mapping.py       | 53 +++++++++++++++++-----------
 2 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/src/geneweaver/client/api/mapping.py b/src/geneweaver/client/api/mapping.py
index 8d3542d..2f6ede0 100644
--- a/src/geneweaver/client/api/mapping.py
+++ b/src/geneweaver/client/api/mapping.py
@@ -30,7 +30,10 @@ def ensembl_mouse_mapping(
     gene_id_type = GeneIdentifier.ENSEMBLE_GENE
 
     if species != Species.MUS_MUSCULUS:
-        gene_id_type = AON_ID_TYPE_FOR_SPECIES[species]
+        try:
+            gene_id_type = AON_ID_TYPE_FOR_SPECIES[species]
+        except KeyError:
+            raise ValueError(f"Species {species} is not supported for ortholog mapping")
 
     response = genesets.get_values(access_token, geneset_id, gene_id_type, in_threshold)
     if species == Species.MUS_MUSCULUS:
diff --git a/tests/unit/api/test_mapping.py b/tests/unit/api/test_mapping.py
index 5d12c80..fdd368a 100644
--- a/tests/unit/api/test_mapping.py
+++ b/tests/unit/api/test_mapping.py
@@ -6,10 +6,15 @@
 from geneweaver.client.api import mapping
 from geneweaver.client.api.aon import OrthologAlgorithms
 from geneweaver.core.enum import Species
+from geneweaver.core.mapping import AON_ID_TYPE_FOR_SPECIES
+
+
+SUPPORTED_AON_SP_IDS = [int(species) for species in AON_ID_TYPE_FOR_SPECIES.keys()]
 
 
 @pytest.mark.parametrize("species_id", [int(species) for species in Species])
 @pytest.mark.parametrize("algorithm", [algorithm for algorithm in OrthologAlgorithms])
+@patch("geneweaver.client.api.mapping.insert_colon_delimiter", lambda x: x)
 @patch("geneweaver.client.api.mapping.genesets.get")
 @patch("geneweaver.client.api.mapping.genesets.get_values")
 @patch("geneweaver.client.api.mapping.aon.algorithm_id_from_name")
@@ -45,28 +50,34 @@ def test_ensembl_mouse_mapping(
 
     mock_aon_map_symbols.side_effect = [{"A1": "1", "B1": "2"}, {"A2": "1", "B2": "2"}]
 
-    result = mapping.ensembl_mouse_mapping(
-        "fake_access_token", 123, True, OrthologAlgorithms.HGNC
-    )
+    if species_id in SUPPORTED_AON_SP_IDS:
+        result = mapping.ensembl_mouse_mapping(
+            "fake_access_token", 123, True, OrthologAlgorithms.HGNC
+        )
+        assert mock_genesets_get.call_count == 1
+        assert mock_genesets_get_values.call_count == 1
 
-    assert mock_genesets_get.call_count == 1
-    assert mock_genesets_get_values.call_count == 1
+        if species_id == 1:
+            assert result == [
+                {"gene_id": "A", "score": "1"},
+                {"gene_id": "B", "score": "2"},
+            ]
+            expected_api_calls, expected_mapping_calls = 0, 0
 
-    if species_id == 1:
-        assert result == [
-            {"gene_id": "A", "score": "1"},
-            {"gene_id": "B", "score": "2"},
-        ]
-        expected_api_calls, expected_mapping_calls = 0, 0
+        else:
+            assert result == [
+                {"gene_id": "A2", "score": "1"},
+                {"gene_id": "B2", "score": "2"},
+            ]
+            expected_api_calls, expected_mapping_calls = 1, 2
 
-    else:
-        assert result == [
-            {"gene_id": "A2", "score": "1"},
-            {"gene_id": "B2", "score": "2"},
-        ]
-        expected_api_calls, expected_mapping_calls = 1, 2
+        assert mock_aon_algorithm_id_from_name.call_count == expected_api_calls
+        assert mock_aon_ortholog_mapping.call_count == expected_api_calls
+        assert mock_genes_mappings.call_count == expected_api_calls
+        assert mock_aon_map_symbols.call_count == expected_mapping_calls
 
-    assert mock_aon_algorithm_id_from_name.call_count == expected_api_calls
-    assert mock_aon_ortholog_mapping.call_count == expected_api_calls
-    assert mock_genes_mappings.call_count == expected_api_calls
-    assert mock_aon_map_symbols.call_count == expected_mapping_calls
+    else:
+        with pytest.raises(ValueError):
+            mapping.ensembl_mouse_mapping(
+                "fake_access_token", 123, True, OrthologAlgorithms.HGNC
+            )