fix: attribute check for tars injection materials (#234)

AllenNeuralDynamics · May 17, 2024 · ae872dc · ae872dc
1 parent ff3bc46
commit ae872dc
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 7 deletions.
diff --git a/src/aind_metadata_service/tars/mapping.py b/src/aind_metadata_service/tars/mapping.py
@@ -218,8 +218,8 @@ def get_virus_strains(response: ModelResponse) -> List:
                 for procedure in subject_procedure.procedures:
                     if (
                         isinstance(procedure, Injection)
-                        and procedure.injection_materials
-                        and procedure.injection_materials[0].name
+                        and hasattr(procedure, "injection_materials")
+                        and hasattr(procedure.injection_materials[0], "name")
                     ):
                         virus_strain = procedure.injection_materials[0].name
                         viruses.append(virus_strain)
@@ -240,14 +240,15 @@ def integrate_injection_materials(
             pre_procedures = response.aind_models[0]
             for subject_procedure in pre_procedures.subject_procedures:
                 for procedure in subject_procedure.procedures:
-                    if (
-                        isinstance(procedure, Injection)
-                        and procedure.injection_materials
+                    if isinstance(procedure, Injection) and hasattr(
+                        procedure, "injection_materials"
                     ):
                         for idx, injection_material in enumerate(
                             procedure.injection_materials
                         ):
-                            if isinstance(injection_material, ViralMaterial):
+                            if isinstance(
+                                injection_material, ViralMaterial
+                            ) and hasattr(injection_material, "name"):
                                 virus_strain = injection_material.name
                                 tars_response = tars_mapping.get(virus_strain)
                                 if (

diff --git a/tests/tars/test_mapping.py b/tests/tars/test_mapping.py
@@ -41,6 +41,16 @@ class TestTarsResponseHandler(unittest.TestCase):
         status_code=StatusCodes.DB_RESPONDED,
     )
 
+    surgery1 = Surgery.model_construct(
+        procedures=[NanojectInjection.model_construct()]
+    )
+    procedures_response1 = ModelResponse(
+        aind_models=[
+            Procedures(subject_id="12345", subject_procedures=[surgery1])
+        ],
+        status_code=StatusCodes.DB_RESPONDED,
+    )
+
     def test_map_prep_type_and_protocol(self):
         """Tests that prep_type and protocol are mapped as expected."""
         (
@@ -283,6 +293,11 @@ def test_get_virus_strains(self):
 
         self.assertEqual(virus_strains, expected_virus_strains)
 
+        no_virus_strains = self.handler.get_virus_strains(
+            self.procedures_response1
+        )
+        self.assertEqual(no_virus_strains, [])
+
     def test_integrate_injection_materials(self):
         """Tests that injection materials are integrated into
         procedures response as expected"""

diff --git a/tests/test_response_handler.py b/tests/test_response_handler.py
@@ -216,7 +216,7 @@ def test_no_validation(self):
             content=(
                 {
                     "message": "Valid Request Format. "
-                               "Models have not been validated.",
+                    "Models have not been validated.",
                     "data": model_json,
                 }
             ),