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 e963204) Signed-off-by: Hailong Cui <[email protected]> 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>
opensearch-project · Feb 7, 2024 · 1a514c6 · 1a514c6
1 parent b40255e
commit 1a514c6
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 8 deletions.
diff --git 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
@@ -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(