feat: adds bool to allow for partial rig matches (#335)

src/aind_metadata_service/server.py

@@ -115,16 +115,20 @@ async def retrieve_bergamo_session(job_settings: BergamoJobSettings):
 
 
 @app.get("/instrument/{instrument_id}")
-async def retrieve_instrument(instrument_id):
+async def retrieve_instrument(instrument_id, partial_match=False):
     """Retrieves instrument from slims"""
-    model_response = slims_client.get_instrument_model_response(instrument_id)
+    model_response = slims_client.get_instrument_model_response(
+        instrument_id, partial_match=partial_match
+    )
     return model_response.map_to_json_response(validate=False)
 
 
 @app.get("/rig/{rig_id}")
-async def retrieve_rig(rig_id):
+async def retrieve_rig(rig_id, partial_match=False):
     """Retrieves rig from slims"""
-    model_response = slims_client.get_rig_model_response(rig_id)
+    model_response = slims_client.get_rig_model_response(
+        rig_id, partial_match=partial_match
+    )
     return model_response.map_to_json_response(validate=False)
 
 

src/aind_metadata_service/slims/client.py

@@ -15,6 +15,7 @@
 from pydantic import Extra, Field, SecretStr
 from pydantic_settings import BaseSettings
 from requests.models import Response
+from slims import criteria
 from slims.criteria import equals
 
 from aind_metadata_service.client import StatusCodes
@@ -55,16 +56,29 @@ def _is_json_file(file: Response) -> bool:
         """Checks whether file is a json."""
         return file.headers.get("Content-Type", "") == "application/json"
 
-    def get_instrument_model_response(self, input_id) -> ModelResponse:
+    def get_instrument_model_response(
+        self, input_id: str, partial_match: bool = False
+    ) -> ModelResponse:
         """
         Fetches a response from SLIMS, extracts Instrument models from the
         response and creates a ModelResponse with said models.
         """
         try:
-            inst = self.client.fetch_model(SlimsInstrumentRdrc, name=input_id)
-            attachment = self.client.fetch_attachment(
-                inst, equals("name", input_id)
-            )
+            if partial_match:
+                inst = self.client.fetch_model(
+                    SlimsInstrumentRdrc,
+                    criteria.contains("name", input_id),
+                )
+                attachment = self.client.fetch_attachment(
+                    inst,
+                )
+            else:
+                inst = self.client.fetch_model(
+                    SlimsInstrumentRdrc, name=input_id
+                )
+                attachment = self.client.fetch_attachment(
+                    inst, equals("name", input_id)
+                )
             if attachment:
                 attachment_response = self.client.fetch_attachment_content(
                     attachment
@@ -91,16 +105,29 @@ def get_instrument_model_response(self, input_id) -> ModelResponse:
             logging.error(repr(e))
             return ModelResponse.internal_server_error_response()
 
-    def get_rig_model_response(self, input_id) -> ModelResponse:
+    def get_rig_model_response(
+        self, input_id: str, partial_match: bool = False
+    ) -> ModelResponse:
         """
         Fetches a response from SLIMS, extracts Rig models from the
         response and creates a ModelResponse with said models.
         """
         try:
-            inst = self.client.fetch_model(SlimsInstrumentRdrc, name=input_id)
-            attachment = self.client.fetch_attachment(
-                inst, equals("name", input_id)
-            )
+            if partial_match:
+                inst = self.client.fetch_model(
+                    SlimsInstrumentRdrc,
+                    criteria.contains("name", input_id),
+                )
+                attachment = self.client.fetch_attachment(
+                    inst,
+                )
+            else:
+                inst = self.client.fetch_model(
+                    SlimsInstrumentRdrc, name=input_id
+                )
+                attachment = self.client.fetch_attachment(
+                    inst, equals("name", input_id)
+                )
             if attachment:
                 attachment_response = self.client.fetch_attachment_content(
                     attachment

tests/slims/test_client.py

@@ -98,6 +98,37 @@ def test_get_instrument_model_response_success(self, mock_is_json_file):
                 SlimsInstrumentRdrc, name="test_id"
             )
 
+    @patch(
+        "aind_metadata_service.slims.client.SlimsHandler._is_json_file",
+    )
+    @patch("slims.criteria.contains")
+    def test_get_instrument_model_response_success_partial_match(
+        self, mock_contains: MagicMock, mock_is_json_file: MagicMock
+    ):
+        """Test successful response from get_instrument_model_response when
+        partial_match is set to True"""
+        mock_inst = MagicMock()
+        mock_is_json_file.return_value = True
+        self.mock_client.fetch_model.return_value = mock_inst
+        mock_attachment = MagicMock()
+        self.mock_client.fetch_attachment.return_value = mock_attachment
+        mock_response = MagicMock(spec=Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = {}
+        self.mock_client.fetch_attachment_content.return_value = mock_response
+
+        with patch(
+            "aind_data_schema.core.instrument.Instrument.model_construct"
+        ) as mock_construct:
+            mock_construct.return_value = MagicMock(spec=Instrument)
+            response = self.handler.get_instrument_model_response(
+                "test_id", partial_match=True
+            )
+
+            self.assertEqual(response.status_code, StatusCodes.DB_RESPONDED)
+            mock_construct.assert_called_once_with(**mock_response.json())
+            mock_contains.assert_called_once_with("name", "test_id")
+
     def test_get_instrument_model_response_invalid_response(self):
         """Test response when the content is not valid response."""
         mock_inst = MagicMock()
@@ -193,6 +224,33 @@ def test_get_rig_model_response_success(self):
                 SlimsInstrumentRdrc, name="test_id"
             )
 
+    @patch("slims.criteria.contains")
+    def test_get_rig_model_response_success_partial_match(
+        self, mock_contains: MagicMock
+    ):
+        """Test successful response from get_rig_model_response when
+        partial_match is set to True."""
+        mock_inst = MagicMock()
+        self.mock_client.fetch_model.return_value = mock_inst
+        mock_attachment = MagicMock()
+        self.mock_client.fetch_attachment.return_value = mock_attachment
+        mock_response = MagicMock(spec=Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = {}
+        self.mock_client.fetch_attachment_content.return_value = mock_response
+
+        with patch(
+            "aind_data_schema.core.rig.Rig.model_construct"
+        ) as mock_construct:
+            mock_construct.return_value = MagicMock(spec=Rig)
+            response = self.handler.get_rig_model_response(
+                "test_id", partial_match=True
+            )
+
+            self.assertEqual(response.status_code, StatusCodes.DB_RESPONDED)
+            mock_construct.assert_called_once_with(**mock_response.json())
+            mock_contains.assert_called_once_with("name", "test_id")
+
     def test_get_rig_model_response_not_found(self):
         """Test when SlimsRecordNotFound is raised."""
         self.mock_client.fetch_model.side_effect = SlimsRecordNotFound