From 5821594e42532a6da6c45e8df7e8a1ef3f42dc0a Mon Sep 17 00:00:00 2001
From: Subhobrata Dey <sbcd90@gmail.com>
Date: Fri, 15 Sep 2023 01:19:17 +0000
Subject: [PATCH] fix flaky integration tests (#581)

Signed-off-by: Subhobrata Dey <sbcd90@gmail.com>
(cherry picked from commit 998d0811156c22a2aa9659229f1e91a5ef8adf3b)
---
 .../SecurityAnalyticsPlugin.java              | 19 -----------
 .../securityanalytics/util/RuleIndices.java   |  1 -
 .../SecurityAnalyticsRestTestCase.java        | 33 +++++--------------
 3 files changed, 8 insertions(+), 45 deletions(-)

diff --git a/src/main/java/org/opensearch/securityanalytics/SecurityAnalyticsPlugin.java b/src/main/java/org/opensearch/securityanalytics/SecurityAnalyticsPlugin.java
index 232b2ea97..2c60321df 100644
--- a/src/main/java/org/opensearch/securityanalytics/SecurityAnalyticsPlugin.java
+++ b/src/main/java/org/opensearch/securityanalytics/SecurityAnalyticsPlugin.java
@@ -292,24 +292,5 @@ public void onFailure(Exception e) {
                 log.warn("Failed to initialize LogType config index and builtin log types");
             }
         });
-        // Trigger initialization of prepackaged rules by calling SearchRule API
-        SearchRequest searchRequest = new SearchRequest(Rule.PRE_PACKAGED_RULES_INDEX);
-        searchRequest.source(new SearchSourceBuilder().query(QueryBuilders.matchAllQuery()).size(0));
-        searchRequest.preference(Preference.PRIMARY_FIRST.type());
-        client.execute(
-                SearchRuleAction.INSTANCE,
-                new SearchRuleRequest(true, searchRequest),
-                new ActionListener<>() {
-                    @Override
-                    public void onResponse(SearchResponse searchResponse) {
-                        log.info("Successfully initialized prepackaged rules");
-                    }
-
-                    @Override
-                    public void onFailure(Exception e) {
-                        log.warn("Failed initializing prepackaged rules", e);
-                    }
-                }
-        );
       }
 }
\ No newline at end of file
diff --git a/src/main/java/org/opensearch/securityanalytics/util/RuleIndices.java b/src/main/java/org/opensearch/securityanalytics/util/RuleIndices.java
index ddb1f8860..baaa8da1c 100644
--- a/src/main/java/org/opensearch/securityanalytics/util/RuleIndices.java
+++ b/src/main/java/org/opensearch/securityanalytics/util/RuleIndices.java
@@ -102,7 +102,6 @@ public void initRuleIndex(ActionListener<CreateIndexResponse> actionListener, bo
         if (!ruleIndexExists(isPrepackaged)) {
             Settings indexSettings = Settings.builder()
                     .put("index.hidden", true)
-                    .put("index.auto_expand_replicas", "0-all")
                     .build();
             CreateIndexRequest indexRequest = new CreateIndexRequest(getRuleIndex(isPrepackaged))
                     .mapping(ruleMappings())
diff --git a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
index 3414dfd39..81fd2f0d3 100644
--- a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
+++ b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
@@ -221,7 +221,7 @@ protected List<Map<String, Object>> searchCorrelatedFindings(String findingId, S
     }
 
     @Before
-    void setDebugLogLevel() throws IOException {
+    void setDebugLogLevel() throws IOException, InterruptedException {
         StringEntity se = new StringEntity("{\n" +
                 "                    \"transient\": {\n" +
                 "                        \"logger.org.opensearch.securityanalytics\":\"DEBUG\",\n" +
@@ -477,30 +477,13 @@ protected Response refreshIndex(String index) throws IOException {
 
     @SuppressWarnings("unchecked")
     protected List<String> getRandomPrePackagedRules() throws IOException {
-        String request = "{\n" +
-                "  \"from\": 0\n," +
-                "  \"size\": 2000\n," +
-                "  \"query\": {\n" +
-                "    \"nested\": {\n" +
-                "      \"path\": \"rule\",\n" +
-                "      \"query\": {\n" +
-                "        \"bool\": {\n" +
-                "          \"must\": [\n" +
-                "            { \"match\": {\"rule.category\": \"" + TestHelpers.randomDetectorType().toLowerCase(Locale.ROOT) + "\"}}\n" +
-                "          ]\n" +
-                "        }\n" +
-                "      }\n" +
-                "    }\n" +
-                "  }\n" +
-                "}";
-
-        Response searchResponse = makeRequest(client(), "POST", String.format(Locale.getDefault(), "%s/_search", SecurityAnalyticsPlugin.RULE_BASE_URI), Collections.singletonMap("pre_packaged", "true"),
-                new StringEntity(request), new BasicHeader("Content-Type", "application/json"));
-        Assert.assertEquals("Searching rules failed", RestStatus.OK, restStatus(searchResponse));
-
-        Map<String, Object> responseBody = asMap(searchResponse);
-        List<Map<String, Object>> hits = ((List<Map<String, Object>>) ((Map<String, Object>) responseBody.get("hits")).get("hits"));
-        return hits.stream().map(hit -> hit.get("_id").toString()).collect(Collectors.toList());
+        return List.of(
+                "36a037c4-c228-4866-b6a3-48eb292b9955",
+                "c6e91a02-d771-4a6d-a700-42587e0b1095",
+                "5a919691-7302-437f-8e10-1fe088afa145",
+                "e5a6b256-3e47-40fc-89d2-7a477edd6915",
+                "06724b9a-52fc-11ed-bdc3-0242ac120002"
+        );
     }
 
     protected List<String> createAggregationRules () throws IOException {