From 721c64fc0958c064f85dfb870318de963bf43aad Mon Sep 17 00:00:00 2001
From: Derek Ho <dxho@amazon.com>
Date: Fri, 13 Oct 2023 12:00:03 -0400
Subject: [PATCH] Fix test and jacoco passing

Signed-off-by: Derek Ho <dxho@amazon.com>
---
 .../service/DataSourceServiceImplTest.java    | 18 +++++-
 .../utils/XContentParserUtilsTest.java        | 60 ++++++++++++++-----
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java
index 0448f72db0..03c6de359f 100644
--- a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java
+++ b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java
@@ -18,8 +18,7 @@
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME;
-import static org.opensearch.sql.datasources.utils.XContentParserUtils.DESCRIPTION_FIELD;
-import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
+import static org.opensearch.sql.datasources.utils.XContentParserUtils.*;
 
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
@@ -307,7 +306,20 @@ void testPatchDefaultDataSource() {
   void testPatchDataSourceSuccessCase() {
     // Tests that patch underlying implementation is to call update
     Map<String, Object> dataSourceData =
-        new HashMap<>(Map.of(NAME_FIELD, "testDS", DESCRIPTION_FIELD, "test"));
+        new HashMap<>(
+            Map.of(
+                NAME_FIELD,
+                "testDS",
+                DESCRIPTION_FIELD,
+                "test",
+                CONNECTOR_FIELD,
+                "PROMETHEUS",
+                ALLOWED_ROLES_FIELD,
+                new ArrayList<>(),
+                PROPERTIES_FIELD,
+                Map.of(),
+                RESULT_INDEX_FIELD,
+                ""));
     DataSourceMetadata getData =
         metadata("testDS", DataSourceType.OPENSEARCH, Collections.emptyList(), ImmutableMap.of());
     when(dataSourceMetadataStorage.getDataSourceMetadata("testDS"))
diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java
index 70df3196bb..42d0e0fea6 100644
--- a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java
+++ b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java
@@ -58,26 +58,39 @@ public void testToDataSourceMetadataFromJson() {
   @Test
   public void testToMapFromJson() {
     Map<String, Object> dataSourceData =
-        new HashMap<>(
-            Map.of(
-                NAME_FIELD,
-                "test_DS",
-                DESCRIPTION_FIELD,
-                "test",
-                ALLOWED_ROLES_FIELD,
-                new ArrayList<>(),
-                PROPERTIES_FIELD,
-                Map.of(),
-                CONNECTOR_FIELD,
-                "PROMETHEUS",
-                RESULT_INDEX_FIELD,
-                ""));
+        Map.of(
+            NAME_FIELD,
+            "test_DS",
+            DESCRIPTION_FIELD,
+            "test",
+            ALLOWED_ROLES_FIELD,
+            new ArrayList<>(),
+            PROPERTIES_FIELD,
+            Map.of(),
+            CONNECTOR_FIELD,
+            "PROMETHEUS",
+            RESULT_INDEX_FIELD,
+            "");
+
+    Map<String, Object> dataSourceDataConnectorRemoved =
+        Map.of(
+            NAME_FIELD,
+            "test_DS",
+            DESCRIPTION_FIELD,
+            "test",
+            ALLOWED_ROLES_FIELD,
+            new ArrayList<>(),
+            PROPERTIES_FIELD,
+            Map.of(),
+            RESULT_INDEX_FIELD,
+            "");
+
     Gson gson = new Gson();
     String json = gson.toJson(dataSourceData);
 
     Map<String, Object> parsedData = XContentParserUtils.toMap(json);
 
-    Assertions.assertEquals(parsedData, dataSourceData);
+    Assertions.assertEquals(parsedData, dataSourceDataConnectorRemoved);
     Assertions.assertEquals("test", parsedData.get(DESCRIPTION_FIELD));
   }
 
@@ -151,4 +164,21 @@ public void testToDataSourceMetadataFromJsonUsingUnknownObject() {
             });
     Assertions.assertEquals("Unknown field: test", exception.getMessage());
   }
+
+  @SneakyThrows
+  @Test
+  public void testToMapFromJsonUsingUnknownObject() {
+    HashMap<String, String> hashMap = new HashMap<>();
+    hashMap.put("test", "test");
+    Gson gson = new Gson();
+    String json = gson.toJson(hashMap);
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> {
+              XContentParserUtils.toMap(json);
+            });
+    Assertions.assertEquals("Unknown field: test", exception.getMessage());
+  }
 }