From 2ff58a92c5884f5bb1cb88ba6a2f4a03f4524ff7 Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Fri, 28 Jul 2023 23:46:03 +0200
Subject: [PATCH] Fixing test classes after core refactoring (#237)

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 .../common/BaseNeuralSearchIT.java             | 18 +++++++-----------
 .../processor/TextEmbeddingProcessorIT.java    |  3 +--
 2 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
index 978951a76..d3636e660 100644
--- a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
+++ b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
@@ -97,7 +97,7 @@ protected String uploadModel(String requestBody) throws Exception {
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT))
         );
         Map<String, Object> uploadResJson = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(uploadResponse.getEntity()),
             false
         );
@@ -126,7 +126,7 @@ protected void loadModel(String modelId) throws Exception {
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT))
         );
         Map<String, Object> uploadResJson = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(uploadResponse.getEntity()),
             false
         );
@@ -175,7 +175,7 @@ protected float[] runInference(String modelId, String queryText) {
         );
 
         Map<String, Object> inferenceResJson = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(inferenceResponse.getEntity()),
             false
         );
@@ -205,7 +205,7 @@ protected void createIndexWithConfiguration(String indexName, String indexConfig
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT))
         );
         Map<String, Object> node = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(response.getEntity()),
             false
         );
@@ -229,7 +229,7 @@ protected void createPipelineProcessor(String modelId, String pipelineName) thro
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT))
         );
         Map<String, Object> node = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(pipelineCreateResponse.getEntity()),
             false
         );
@@ -295,7 +295,7 @@ protected Map<String, Object> search(String index, QueryBuilder queryBuilder, Qu
 
         String responseBody = EntityUtils.toString(response.getEntity());
 
-        return XContentHelper.convertToMap(XContentFactory.xContent(XContentType.JSON), responseBody, false);
+        return XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody, false);
     }
 
     /**
@@ -406,11 +406,7 @@ protected Map<String, Object> getTaskQueryResponse(String taskId) throws Excepti
             toHttpEntity(""),
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT))
         );
-        return XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
-            EntityUtils.toString(taskQueryResponse.getEntity()),
-            false
-        );
+        return XContentHelper.convertToMap(XContentType.JSON.xContent(), EntityUtils.toString(taskQueryResponse.getEntity()), false);
     }
 
     protected boolean checkComplete(Map<String, Object> node) {
diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java
index 226a3878d..7cea20a41 100644
--- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java
+++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java
@@ -13,7 +13,6 @@
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.apache.hc.core5.http.message.BasicHeader;
 import org.opensearch.client.Response;
-import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.neuralsearch.common.BaseNeuralSearchIT;
@@ -71,7 +70,7 @@ private void ingestDocument() throws Exception {
             ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
         );
         Map<String, Object> map = XContentHelper.convertToMap(
-            XContentFactory.xContent(XContentType.JSON),
+            XContentType.JSON.xContent(),
             EntityUtils.toString(response.getEntity()),
             false
         );