SageMaker: search() now supports the CONTAINS filter (#8020)

moto/sagemaker/models.py

@@ -11,6 +11,7 @@
 
 from moto.core.base_backend import BackendDict, BaseBackend
 from moto.core.common_models import BaseModel, CloudFormationModel
+from moto.core.utils import camelcase_to_underscores
 from moto.sagemaker import validators
 from moto.utilities.paginator import paginate
 from moto.utilities.utils import ARN_PARTITION_REGEX, get_partition
@@ -3055,6 +3056,9 @@ def list_experiments(self) -> List["FakeExperiment"]:
         return list(self.experiments.values())
 
     def search(self, resource: Any = None, search_expression: Any = None) -> Any:
+        """
+        Only a few SearchExpressions are implemented. Please open a bug report if you find any issues.
+        """
         next_index = None
 
         valid_resources = [
@@ -3078,50 +3082,54 @@ def search(self, resource: Any = None, search_expression: Any = None) -> Any:
                 f"An error occurred (ValidationException) when calling the Search operation: 1 validation error detected: Value '{resource}' at 'resource' failed to satisfy constraint: Member must satisfy enum value set: {valid_resources}"
             )
 
+        def compare_value(actual: Any, expected: Any, operator: str) -> bool:
+            # Defeault:  operator == "Equals"
+            if operator == "Contains":
+                return expected in actual
+            if operator == "NotEquals":
+                return expected != actual
+            return actual == expected
+
         def evaluate_search_expression(item: Any) -> bool:
             filters = None
             if search_expression is not None:
                 filters = search_expression.get("Filters")
 
             if filters is not None:
                 for f in filters:
-                    if f["Operator"] == "Equals":
-                        if f["Name"].startswith("Tags."):
-                            key = f["Name"][5:]
-                            value = f["Value"]
-
-                            if (
-                                len(
-                                    [
-                                        e
-                                        for e in item.tags
-                                        if e["Key"] == key and e["Value"] == value
-                                    ]
-                                )
-                                == 0
-                            ):
-                                return False
-                        if f["Name"] == "ExperimentName":
-                            experiment_name = f["Value"]
-
-                            if hasattr(item, "experiment_name"):
-                                if getattr(item, "experiment_name") != experiment_name:
-                                    return False
-                            else:
-                                raise ValidationError(
-                                    message="Unknown property name: ExperimentName"
-                                )
-
-                        if f["Name"] == "TrialName":
-                            raise AWSValidationException(
-                                f"An error occurred (ValidationException) when calling the Search operation: Unknown property name: {f['Name']}"
-                            )
-
-                        if f["Name"] == "Parents.TrialName":
-                            trial_name = f["Value"]
-
-                            if getattr(item, "trial_name") != trial_name:
+                    prop_key = camelcase_to_underscores(f["Name"])
+                    if f["Name"].startswith("Tags."):
+                        key = f["Name"][5:]
+                        value = f["Value"]
+
+                        if f["Operator"] == "Equals":
+                            if not [
+                                e
+                                for e in item.tags
+                                if e["Key"] == key and e["Value"] == value
+                            ]:
                                 return False
+                        return True
+
+                    elif f["Name"] == "TrialName":
+                        raise AWSValidationException(
+                            f"An error occurred (ValidationException) when calling the Search operation: Unknown property name: {f['Name']}"
+                        )
+
+                    elif f["Name"] == "Parents.TrialName":
+                        trial_name = f["Value"]
+                        if getattr(item, "trial_name") != trial_name:
+                            return False
+
+                    elif hasattr(item, prop_key):
+                        if not compare_value(
+                            getattr(item, prop_key), f["Value"], f["Operator"]
+                        ):
+                            return False
+                    else:
+                        raise ValidationError(
+                            message=f"Unknown property name: {f['Name']}"
+                        )
 
             return True
 

tests/test_sagemaker/test_sagemaker_search.py

@@ -100,6 +100,61 @@ def test_search_trial_component_with_experiment_name(sagemaker_client):
     assert ex.value.response["ResponseMetadata"]["HTTPStatusCode"] == 400
 
 
+def test_search_experiments_with_filter(sagemaker_client):
+    experiment_name = "experiment_name"
+    trial_component_name = "trial_component_name"
+    _set_up_trial_component(
+        sagemaker_client,
+        experiment_name=experiment_name,
+        trial_component_name=trial_component_name,
+    )
+
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "Contains", "Value": "_DEV"}, 0
+    )
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "Contains", "Value": "_name"}, 1
+    )
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "Equals", "Value": "_name"}, 0
+    )
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "Equals", "Value": experiment_name}, 1
+    )
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "NotEquals", "Value": experiment_name}, 0
+    )
+    _verify_search_results(
+        {"Name": "ExperimentName", "Operator": "NotEquals", "Value": "else"}, 1
+    )
+
+
+def test_search_model_package_groups_with_filter(sagemaker_client):
+    sagemaker_client.create_model_package_group(
+        ModelPackageGroupName="mpg_DEV",
+        ModelPackageGroupDescription="test-model-package-group-description",
+    )
+
+    _verify_search_results(
+        {"Name": "ModelPackageGroupName", "Operator": "Contains", "Value": "_DEV"},
+        1,
+        resource="ModelPackageGroup",
+    )
+    _verify_search_results(
+        {"Name": "ModelPackageGroupName", "Operator": "Contains", "Value": "_PROD"},
+        0,
+        resource="ModelPackageGroup",
+    )
+
+
+def _verify_search_results(_filter, nr_of_results, resource="Experiment"):
+    sagemaker_client = boto3.client("sagemaker", region_name=TEST_REGION_NAME)
+    results = sagemaker_client.search(
+        Resource=resource, SearchExpression={"Filters": [_filter]}
+    )["Results"]
+    assert len(results) == nr_of_results
+
+
 def _set_up_trial_component(
     sagemaker_client,
     experiment_name="some-experiment-name",