Updated pr.

Author: lu-ohai

tests/unitary/with_extras/langchain/embeddings/test_oci_model_deployment_endpoint.py

@@ -6,30 +6,49 @@
 
 """Test OCI Data Science Model Deployment Endpoint."""
 
-import responses
-from pytest_mock import MockerFixture
+from unittest.mock import MagicMock, patch
 from ads.llm import OCIModelDeploymentEndpointEmbeddings
 
 
-@responses.activate
-def test_embedding_call(mocker: MockerFixture) -> None:
+@patch("ads.llm.OCIModelDeploymentEndpointEmbeddings._embed_with_retry")
+def test_embed_documents(mock_embed_with_retry) -> None:
     """Test valid call to oci model deployment endpoint."""
-    endpoint = "https://MD_OCID/predict"
-    documents = ["Hello", "World"]
     expected_output = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
-    responses.add(
-        responses.POST,
-        endpoint,
-        json={
+    result = MagicMock()
+    result.json = MagicMock(
+        return_value={
             "data": [{"embedding": expected_output}],
-        },
-        status=200,
+        }
     )
-    mocker.patch("ads.common.auth.default_signer", return_value=dict(signer=None))
+    mock_embed_with_retry.return_value = result
+    endpoint = "https://MD_OCID/predict"
+    documents = ["Hello", "World"]
 
     embeddings = OCIModelDeploymentEndpointEmbeddings(
         endpoint=endpoint,
     )
 
     output = embeddings.embed_documents(documents)
     assert output == expected_output
+
+
+@patch("ads.llm.OCIModelDeploymentEndpointEmbeddings._embed_with_retry")
+def test_embed_query(mock_embed_with_retry) -> None:
+    """Test valid call to oci model deployment endpoint."""
+    expected_output = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
+    result = MagicMock()
+    result.json = MagicMock(
+        return_value={
+            "data": [{"embedding": expected_output}],
+        }
+    )
+    mock_embed_with_retry.return_value = result
+    endpoint = "https://MD_OCID/predict"
+    query = "Hello world"
+
+    embeddings = OCIModelDeploymentEndpointEmbeddings(
+        endpoint=endpoint,
+    )
+
+    output = embeddings.embed_query(query)
+    assert output == expected_output[0]