Add test for map getter in clients

Signed-off-by: Daniel Widdis <[email protected]>
dbwiddis · Jun 14, 2024 · e63e98f · e63e98f
1 parent 411550d
commit e63e98f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/plugin/src/test/java/org/opensearch/ml/sdkclient/LocalClusterIndicesClientTests.java b/plugin/src/test/java/org/opensearch/ml/sdkclient/LocalClusterIndicesClientTests.java
@@ -43,6 +43,8 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.concurrent.OpenSearchExecutors;
+import org.opensearch.common.xcontent.XContentHelper;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
@@ -143,7 +145,9 @@ public void testGetDataObject() throws IOException {
         GetResponse getResponse = mock(GetResponse.class);
         when(getResponse.isExists()).thenReturn(true);
         when(getResponse.getId()).thenReturn(TEST_ID);
-        when(getResponse.getSourceAsString()).thenReturn(testDataObject.toJson());
+        String json = testDataObject.toJson();
+        when(getResponse.getSourceAsString()).thenReturn(json);
+        when(getResponse.getSource()).thenReturn(XContentHelper.convertToMap(JsonXContent.jsonXContent, json, false));
         @SuppressWarnings("unchecked")
         ActionFuture<GetResponse> future = mock(ActionFuture.class);
         when(mockedClient.get(any(GetRequest.class))).thenReturn(future);
@@ -158,6 +162,7 @@ public void testGetDataObject() throws IOException {
         verify(mockedClient, times(1)).get(requestCaptor.capture());
         assertEquals(TEST_INDEX, requestCaptor.getValue().index());
         assertEquals(TEST_ID, response.id());
+        assertEquals("foo", response.source().get("data"));
         assertTrue(response.parser().isPresent());
         XContentParser parser = response.parser().get();
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);

diff --git a/plugin/src/test/java/org/opensearch/ml/sdkclient/RemoteClusterIndicesClientTests.java b/plugin/src/test/java/org/opensearch/ml/sdkclient/RemoteClusterIndicesClientTests.java
@@ -173,6 +173,7 @@ public void testGetDataObject() throws IOException {
 
         assertEquals(TEST_INDEX, getRequestCaptor.getValue().index());
         assertEquals(TEST_ID, response.id());
+        assertEquals("foo", response.source().get("data"));
         assertTrue(response.parser().isPresent());
         XContentParser parser = response.parser().get();
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);