Additional test case added for each search api: task and model

Signed-off-by: Nurlan <nabzalbekov0@gmail.com>
opensearch-project · Jun 16, 2023 · 867e708 · 867e708
1 parent 33be035
commit 867e708
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/opensearch_py_ml/ml_commons/ml_commons_client.py b/opensearch_py_ml/ml_commons/ml_commons_client.py
@@ -391,6 +391,8 @@ def search_task(self, input_json) -> object:
         if isinstance(input_json, str):
             try:
                 json_obj = json.loads(input_json)
+                if not isinstance(json_obj, dict):
+                    return "Invalid JSON object passed as argument."
                 API_BODY = json.dumps(json_obj)
             except json.JSONDecodeError:
                 return "Invalid JSON string passed as argument."
@@ -419,6 +421,8 @@ def search_model(self, input_json) -> object:
         if isinstance(input_json, str):
             try:
                 json_obj = json.loads(input_json)
+                if not isinstance(json_obj, dict):
+                    return "Invalid JSON object passed as argument."
                 API_BODY = json.dumps(json_obj)
             except json.JSONDecodeError:
                 return "Invalid JSON string passed as argument."

diff --git a/tests/common.py b/tests/common.py
@@ -57,7 +57,7 @@
 )
 _pd_ecommerce.insert(2, "customer_birth_date", None)
 _pd_ecommerce.index = _pd_ecommerce.index.map(str)  # make index 'object' not int
-_pd_ecommerce["customer_birth_date"].astype("datetime64[ns]")
+_pd_ecommerce["customer_birth_date"].astype("datetime64")
 _oml_ecommerce = oml.DataFrame(OPENSEARCH_TEST_CLIENT, ECOMMERCE_INDEX_NAME)
 
 

diff --git a/tests/ml_commons/test_ml_commons_client.py b/tests/ml_commons/test_ml_commons_client.py
@@ -130,6 +130,24 @@ def test_search():
         raised = True
     assert raised == False, "Raised Exception in searching task"
 
+    raised = False
+    try:
+        search_task_obj = ml_client.search_task(input_json="15")
+        assert search_task_obj == "Invalid JSON object passed as argument."
+    except:  # noqa: E722
+        raised = True
+    assert raised == False, "Raised Exception in searching task"
+
+    raised = False
+    try:
+        search_task_obj = ml_client.search_task(
+            input_json='{"query": {"match_all": {}},size: 1}'
+        )
+        assert search_task_obj == "Invalid JSON string passed as argument."
+    except:  # noqa: E722
+        raised = True
+    assert raised == False, "Raised Exception in searching task"
+
     # Search model cases
     raised = False
     try:
@@ -159,6 +177,24 @@ def test_search():
         raised = True
     assert raised == False, "Raised Exception in searching model"
 
+    raised = False
+    try:
+        search_model_obj = ml_client.search_model(input_json="15")
+        assert search_model_obj == "Invalid JSON object passed as argument."
+    except:  # noqa: E722
+        raised = True
+    assert raised == False, "Raised Exception in searching model"
+
+    raised = False
+    try:
+        search_model_obj = ml_client.search_model(
+            input_json='{"query": {"match_all": {}},size: 1}'
+        )
+        assert search_model_obj == "Invalid JSON string passed as argument."
+    except:  # noqa: E722
+        raised = True
+    assert raised == False, "Raised Exception in searching model"
+
 
 def test_DEPRECATED_integration_pretrained_model_upload_unload_delete():
     raised = False