From 42372a01185cd6b77ded4e4cbd3b901afb39bf84 Mon Sep 17 00:00:00 2001
From: Surya Sashank Nistala <snistala@amazon.com>
Date: Fri, 31 May 2024 19:22:02 -0700
Subject: [PATCH] test failing due to node closed

---
 .../iocscan/dao/IocMatchService.java          |  4 +-
 .../SecurityAnalyticsIntegTestCase.java       |  2 -
 .../iocscan/dao/IocMatchServiceIT.java        | 44 ++++++++++++-------
 3 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/src/main/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchService.java b/src/main/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchService.java
index 2bbe0c9f6..9bf92a985 100644
--- a/src/main/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchService.java
+++ b/src/main/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchService.java
@@ -133,7 +133,7 @@ public void createIndexIfNotExists(final ActionListener<Void> listener) {
             }
             final CreateIndexRequest createIndexRequest = new CreateIndexRequest(INDEX_NAME).mapping(getIndexMapping())
                     .settings(SecurityAnalyticsPlugin.TIF_JOB_INDEX_SETTING);
-            StashedThreadContext.run(client, () -> client.admin().indices().create(createIndexRequest, ActionListener.wrap(
+            client.admin().indices().create(createIndexRequest, ActionListener.wrap(
                     r -> {
                         log.debug("Ioc match index created");
                         listener.onResponse(null);
@@ -146,7 +146,7 @@ public void createIndexIfNotExists(final ActionListener<Void> listener) {
                         log.error("Failed to create security analytics threat intel job index", e);
                         listener.onFailure(e);
                     }
-            )));
+            ));
         } catch (Exception e) {
             log.error("Failure in creating ioc_match index", e);
             listener.onFailure(e);
diff --git a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsIntegTestCase.java b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsIntegTestCase.java
index 14a6396d2..a85196663 100644
--- a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsIntegTestCase.java
+++ b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsIntegTestCase.java
@@ -15,8 +15,6 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
                 );
     }
 
-    settings
-
     @Override
     protected boolean ignoreExternalCluster() {
         return true;
diff --git a/src/test/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchServiceIT.java b/src/test/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchServiceIT.java
index f28ad32cb..fc24b5b76 100644
--- a/src/test/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchServiceIT.java
+++ b/src/test/java/org/opensearch/securityanalytics/threatIntel/iocscan/dao/IocMatchServiceIT.java
@@ -1,11 +1,12 @@
 package org.opensearch.securityanalytics.threatIntel.iocscan.dao;
 
 import org.opensearch.action.LatchedActionListener;
+import org.opensearch.action.StepListener;
 import org.opensearch.action.search.SearchRequest;
+import org.opensearch.common.util.concurrent.CountDown;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.securityanalytics.SecurityAnalyticsIntegTestCase;
 import org.opensearch.securityanalytics.model.threatintel.IocMatch;
-import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.time.Instant;
 import java.util.ArrayList;
@@ -15,22 +16,33 @@
 
 public class IocMatchServiceIT extends SecurityAnalyticsIntegTestCase {
 
-    public void test_indexIocMatches() {
+    public void test_indexIocMatches() throws InterruptedException {
         IocMatchService service = new IocMatchService(client(), clusterService());
         List<IocMatch> iocMatches = generateIocMatches(10);
-        CountDownLatch latch = new CountDownLatch(1);
-        service.indexIocMatches(iocMatches, new LatchedActionListener<>(new ActionListener<>() {
-            @Override
-            public void onResponse(Void unused) {
-                client().search(new SearchRequest(IocMatchService.INDEX_NAME)).
-            }
+        CountDown countdown = new CountDown(1);
+        service.indexIocMatches(iocMatches, ActionListener.wrap(r -> {
+            countdown.countDown();
+        }, e -> {
+            logger.error("failed to index ioc matches", e);
+            fail();
+            countdown.countDown();
+        }));
+        SearchRequest request = new SearchRequest(IocMatchService.INDEX_NAME);
+        request.source().size(10);
+        CountDown countDownLatch1 = new CountDown(1);
+        client().search(request, ActionListener.wrap(
+                response -> {
+                    assertEquals(response.getHits().getHits().length, 10);
+                    countDownLatch1.countDown();
+                },
+                e -> {
+                    logger.error("failed to search indexed ioc matches", e);
+                    fail();
+                    countDownLatch1.countDown();
+                }
 
-            @Override
-            public void onFailure(Exception e) {
-                logger.error("failed to index ioc matches", e);
-                fail();
-            }
-        }, latch));
+        ));
+        countDownLatch1.isCountedDown();
     }
 
     private List<IocMatch> generateIocMatches(int i) {
@@ -40,8 +52,8 @@ private List<IocMatch> generateIocMatches(int i) {
         for (int i1 = 0; i1 < i; i1++) {
             iocMatches.add(new IocMatch(
                     randomAlphaOfLength(10),
-                    randomList(1,10, () -> randomAlphaOfLength(10)),//docids
-                    randomList(1,10, () -> randomAlphaOfLength(10)), //feedids
+                    randomList(1, 10, () -> randomAlphaOfLength(10)),//docids
+                    randomList(1, 10, () -> randomAlphaOfLength(10)), //feedids
                     monitorId,
                     monitorName,
                     randomAlphaOfLength(10),