Fixes source mapping (#227)

* fix source, linters * Other instead of None
AllenNeuralDynamics · May 10, 2024 · 00cc766 · 00cc766
1 parent 84815d0
commit 00cc766
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 24 deletions.
diff --git a/src/aind_metadata_service/labtracks/client.py b/src/aind_metadata_service/labtracks/client.py
@@ -416,15 +416,26 @@ def map_response_to_subject(self, results: List[dict]) -> List[Subject]:
                 room_id=result.get(SubjectQueryColumns.ROOM_ID.value),
                 cage_id=result.get(SubjectQueryColumns.CAGE_ID.value),
             )
-            breeding_info = BreedingInfo.model_construct(
-                breeding_group=breeding_group,
-                paternal_genotype=paternal_genotype,
-                paternal_id=paternal_id_str,
-                maternal_genotype=maternal_genotype,
-                maternal_id=maternal_id_str,
-            )
-            # Assume all mice from LabTracks is AI?
-            source = Organization.AI
+            breeding_values = [
+                breeding_group,
+                paternal_genotype,
+                paternal_id_str,
+                maternal_genotype,
+                maternal_id_str,
+            ]
+            if any(value is not None for value in breeding_values):
+                breeding_info = BreedingInfo.model_construct(
+                    breeding_group=breeding_group,
+                    paternal_genotype=paternal_genotype,
+                    paternal_id=paternal_id_str,
+                    maternal_genotype=maternal_genotype,
+                    maternal_id=maternal_id_str,
+                )
+                # If breeding info, subject is from AI
+                source = Organization.AI
+            else:
+                breeding_info = None
+                source = Organization.OTHER
             subject = Subject.model_construct(
                 source=source,
                 subject_id=subject_id_str,

diff --git a/src/aind_metadata_service/server.py b/src/aind_metadata_service/server.py
@@ -295,9 +295,9 @@ async def retrieve_procedures(subject_id, pickle: bool = False):
         # smartsheet_response = await retrieve_protocols(
         #     protocol_name=protocol_name
         # )
-        protocols_mapping[protocol_name] = (
-            model_response.map_to_json_response()
-        )
+        protocols_mapping[
+            protocol_name
+        ] = model_response.map_to_json_response()
     integrated_response = protocols_integrator.integrate_protocols(
         response=integrated_response, protocols_mapping=protocols_mapping
     )

diff --git a/src/aind_metadata_service/sharepoint/client.py b/src/aind_metadata_service/sharepoint/client.py
@@ -240,9 +240,9 @@ def _merge_two_responses(
             )
         else:
             left_procedures: List[Procedures] = left_model_response.aind_models
-            right_procedures: List[Procedures] = (
-                right_model_response.aind_models
-            )
+            right_procedures: List[
+                Procedures
+            ] = right_model_response.aind_models
             procedures = self._merge_procedures(
                 left_procedures=left_procedures,
                 right_procedures=right_procedures,

diff --git a/src/aind_metadata_service/smartsheet/protocols/mapping.py b/src/aind_metadata_service/smartsheet/protocols/mapping.py
@@ -91,10 +91,10 @@ def _get_protocol_list(
         """
         rows_associated_with_protocol_name: List[ProtocolInformation] = []
         for row in self.model.rows:
-            opt_protocol: Optional[ProtocolInformation] = (
-                self._map_row_to_protocol(
-                    row=row, input_protocol_name=protocol_name
-                )
+            opt_protocol: Optional[
+                ProtocolInformation
+            ] = self._map_row_to_protocol(
+                row=row, input_protocol_name=protocol_name
             )
             if opt_protocol is not None:
                 rows_associated_with_protocol_name.append(opt_protocol)

diff --git a/src/aind_metadata_service/tars/mapping.py b/src/aind_metadata_service/tars/mapping.py
@@ -265,9 +265,9 @@ def integrate_injection_materials(
                                     new_material.titer = (
                                         injection_material.titer
                                     )
-                                    procedure.injection_materials[idx] = (
-                                        new_material
-                                    )
+                                    procedure.injection_materials[
+                                        idx
+                                    ] = new_material
                                 elif (
                                     tars_response.status_code
                                     == StatusCodes.NO_DATA_FOUND.value

diff --git a/tests/labtracks/test_client.py b/tests/labtracks/test_client.py
@@ -285,7 +285,23 @@ def fetchall():
                         "C57BL/6J",
                         "1234",
                         "000",
-                    ]
+                    ],
+                    [
+                        decimal.Decimal("115977"),
+                        TestResponseExamples.class_values_str,
+                        "M",
+                        datetime.datetime(2012, 5, 13, 0, 0),
+                        None,
+                        None,
+                        None,
+                        None,
+                        None,
+                        "mouse",
+                        None,
+                        "C57BL/6J",
+                        "1234",
+                        "000",
+                    ],
                 ]
 
             @staticmethod
@@ -299,7 +315,10 @@ def close():
 
         subject_response = self.lb_client.get_subject_info(subject_id)
         expected_response = ModelResponse(
-            aind_models=[TestResponseExamples.expected_subject],
+            aind_models=[
+                TestResponseExamples.expected_subject,
+                TestResponseExamples.expected_subject_2,
+            ],
             status_code=StatusCodes.DB_RESPONDED,
         )
         self.assertEqual(

diff --git a/tests/labtracks/test_response_handler.py b/tests/labtracks/test_response_handler.py
@@ -102,6 +102,23 @@ class TestResponseExamples:
         notes=None,
     )
 
+    expected_subject_2 = Subject(
+        species=Species.MUS_MUSCULUS,
+        breeding_info=None,
+        subject_id="115977",
+        sex=Sex.MALE,
+        source=Organization.OTHER,
+        date_of_birth=date(2012, 5, 13),
+        genotype="Adora2a-Cre/wt",
+        alleles=[],
+        background_strain=BackgroundStrain.C57BL_6J,
+        housing=expected_housing,
+        rrid=None,
+        restrictions=None,
+        wellness_reports=[],
+        notes=None,
+    )
+
     test_procedures_response = [
         {
             "id": Decimal(00000),