try to fix security test

Signed-off-by: xinyual <[email protected]>
opensearch-project · Feb 1, 2025 · c83c4f7 · c83c4f7
1 parent 1e82c0b
commit c83c4f7
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java b/src/test/java/org/opensearch/integTest/ToolIntegrationTest.java
@@ -44,6 +44,8 @@ public void setupTestAgent() throws IOException, InterruptedException {
         server.start();
         clusterSettings(false);
         connectorId = setUpConnectorWithRetry(5);
+        TimeUnit.SECONDS.sleep(3);
+        queryConnector();
         modelGroupId = setupModelGroup();
         modelId = setupLLMModel(connectorId, modelGroupId);
         // wait for model to get deployed
@@ -64,10 +66,12 @@ public void stopMockLLM() {
 
     @After
     public void deleteModel() {
-        deleteModel(modelId);
+        if (modelId != null) {
+            deleteModel(modelId);
+        }
     }
 
-    private String setUpConnectorWithRetry(int maxRetryTimes) throws InterruptedException {
+    private String setUpConnectorWithRetry(int maxRetryTimes) throws InterruptedException, IOException {
         int retryTimes = 0;
         String connectorId = null;
         while (retryTimes < maxRetryTimes) {
@@ -81,6 +85,7 @@ private String setUpConnectorWithRetry(int maxRetryTimes) throws InterruptedExce
                 TimeUnit.SECONDS.sleep(20);
             }
         }
+        queryConnector();
         return connectorId;
     }
 
@@ -155,6 +160,16 @@ private String setupModelGroup() throws IOException {
         return JsonParser.parseString(resp).getAsJsonObject().get("model_group_id").getAsString();
     }
 
+    private String queryConnector() throws IOException {
+        Request request = new Request("POST", "/.plugins-ml-connector/_search");
+        request.setJsonEntity("{\n" + "    \"query\": {\n" + "      \"match_all\": {}\n" + "} \n" + "} \n");
+        Response response = executeRequest(request);
+
+        String resp = readResponse(response);
+
+        return JsonParser.parseString(resp).getAsJsonObject().getAsString();
+    }
+
     private String setupLLMModel(String connectorId, String modelGroupId) throws IOException {
         Request request = new Request("POST", "/_plugins/_ml/models/_register?deploy=true");
         request