From 1a514c6c5ef2627f0dee5716f1001e1f83b0c890 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Wed, 7 Feb 2024 09:49:59 +0800
Subject: [PATCH] Increase wait time to avoid flaky test (#173) (#176)

* increae wait time to avoid flaky test



* add retry for setinig up connector



* fix size hardcoded



---------


(cherry picked from commit e9632040e47b24963a86b38a68d723d530033eeb)

Signed-off-by: Hailong Cui <ihailong@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../agent/tools/VisualizationsTool.java       |  2 +-
 .../integTest/ToolIntegrationTest.java        | 25 +++++++++++++------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/main/java/org/opensearch/agent/tools/VisualizationsTool.java b/src/main/java/org/opensearch/agent/tools/VisualizationsTool.java
index 35bb5c1c..2fa6b996 100644
--- a/src/main/java/org/opensearch/agent/tools/VisualizationsTool.java
+++ b/src/main/java/org/opensearch/agent/tools/VisualizationsTool.java
@@ -77,7 +77,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
         boolQueryBuilder.must().add(QueryBuilders.matchQuery(SAVED_OBJECT_TYPE + ".title", parameters.get("input")));
 
         SearchSourceBuilder searchSourceBuilder = SearchSourceBuilder.searchSource().query(boolQueryBuilder);
-        searchSourceBuilder.from(0).size(3);
+        searchSourceBuilder.from(0).size(size);
         SearchRequest searchRequest = Requests.searchRequest(index).source(searchSourceBuilder);
 
         client.search(searchRequest, new ActionListener<>() {
diff --git a/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java b/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java
index 74c3d5cf..315812b1 100644
--- a/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java
+++ b/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java
@@ -43,13 +43,7 @@ public void setupTestAgent() throws IOException, InterruptedException {
         server = MockHttpServer.setupMockLLM(promptHandlers());
         server.start();
         clusterSettings(false);
-        try {
-            connectorId = setUpConnector();
-        } catch (Exception e) {
-            // Wait for ML encryption master key has been initialized
-            TimeUnit.SECONDS.sleep(10);
-            connectorId = setUpConnector();
-        }
+        connectorId = setUpConnectorWithRetry(5);
         modelGroupId = setupModelGroup();
         modelId = setupLLMModel(connectorId, modelGroupId);
         // wait for model to get deployed
@@ -73,6 +67,23 @@ public void deleteModel() {
         deleteModel(modelId);
     }
 
+    private String setUpConnectorWithRetry(int maxRetryTimes) throws InterruptedException {
+        int retryTimes = 0;
+        String connectorId = null;
+        while (retryTimes < maxRetryTimes) {
+            try {
+                connectorId = setUpConnector();
+                break;
+            } catch (Exception e) {
+                // Wait for ML encryption master key has been initialized
+                log.info("Failed to setup connector, retry times: {}", retryTimes);
+                retryTimes++;
+                TimeUnit.SECONDS.sleep(10);
+            }
+        }
+        return connectorId;
+    }
+
     private String setUpConnector() {
         String url = String.format(Locale.ROOT, "http://127.0.0.1:%d/invoke", server.getAddress().getPort());
         return createConnector(