From d96899e1b03d2da0df49b83873ade3ec8e681f92 Mon Sep 17 00:00:00 2001
From: ulixius9 <mayursingal9@gmail.com>
Date: Fri, 26 Jul 2024 00:03:15 +0530
Subject: [PATCH] fix kpi test

---
 .../resources/kpi/KpiResourceTest.java        | 36 ++++++++++++++++---
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/kpi/KpiResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/kpi/KpiResourceTest.java
index ff2d6eb3e521..80fe51e7b40b 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/kpi/KpiResourceTest.java
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/kpi/KpiResourceTest.java
@@ -3,6 +3,7 @@
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.openmetadata.service.Entity.getSearchRepository;
 import static org.openmetadata.service.security.SecurityUtil.getPrincipalName;
 import static org.openmetadata.service.util.EntityUtil.fieldUpdated;
 import static org.openmetadata.service.util.TestUtils.ADMIN_AUTH_HEADERS;
@@ -11,6 +12,7 @@
 import static org.openmetadata.service.util.TestUtils.assertListNull;
 import static org.openmetadata.service.util.TestUtils.assertResponse;
 
+import es.org.elasticsearch.client.RestClient;
 import java.io.IOException;
 import java.util.Map;
 import java.util.UUID;
@@ -23,11 +25,14 @@
 import org.openmetadata.schema.api.dataInsight.kpi.CreateKpiRequest;
 import org.openmetadata.schema.api.dataInsight.kpi.KpiDataInsightChart;
 import org.openmetadata.schema.dataInsight.kpi.Kpi;
-import org.openmetadata.schema.dataInsight.type.KpiResult;
 import org.openmetadata.schema.dataInsight.type.KpiTarget;
 import org.openmetadata.schema.dataInsight.type.KpiTargetType;
+import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
 import org.openmetadata.schema.type.ChangeDescription;
 import org.openmetadata.service.Entity;
+import org.openmetadata.service.apps.bundles.insights.search.DataInsightsSearchInterface;
+import org.openmetadata.service.apps.bundles.insights.search.elasticsearch.ElasticSearchDataInsightsClient;
+import org.openmetadata.service.apps.bundles.insights.search.opensearch.OpenSearchDataInsightsClient;
 import org.openmetadata.service.resources.EntityResourceTest;
 import org.openmetadata.service.util.JsonUtils;
 
@@ -39,7 +44,32 @@ public KpiResourceTest() {
     supportsPatch = false;
   }
 
+  private void createDataAssetsDataStream() {
+    DataInsightsSearchInterface searchInterface;
+    if (getSearchRepository()
+        .getSearchType()
+        .equals(ElasticSearchConfiguration.SearchType.ELASTICSEARCH)) {
+      searchInterface =
+          new ElasticSearchDataInsightsClient(
+              (RestClient) getSearchRepository().getSearchClient().getLowLevelClient());
+    } else {
+      searchInterface =
+          new OpenSearchDataInsightsClient(
+              (os.org.opensearch.client.RestClient)
+                  getSearchRepository().getSearchClient().getLowLevelClient());
+    }
+
+    try {
+      if (!searchInterface.dataAssetDataStreamExists("di-data-assets")) {
+        searchInterface.createDataAssetsDataStream();
+      }
+    } catch (IOException ex) {
+      LOG.error("Couldn't install DataInsightsApp: Can't initialize ElasticSearch Index.");
+    }
+  }
+
   public void setupKpi() throws IOException {
+    createDataAssetsDataStream();
     KPI_TARGET = new KpiTarget().withName("Percentage").withValue("80.0");
   }
 
@@ -67,10 +97,6 @@ void createUpdate_tests_200() throws IOException {
     validateCreatedEntity(createdKpi, create, ADMIN_AUTH_HEADERS);
   }
 
-  private void verifyKpiResult(KpiResult expected, KpiResult actual) {
-    assertEquals(expected, actual);
-  }
-
   @Override
   public CreateKpiRequest createRequest(String name) {
     try {