Fix failing test

rithin-pullela-aws · rithin-pullela-aws · commit cbeed7aa1913 · 2025-08-05T12:44:31.000-07:00
Signed-off-by: rithin-pullela-aws &lt;rithinp@amazon.com&gt;
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestQueryPlanningToolIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestQueryPlanningToolIT.java
@@ -74,6 +74,8 @@ public void setup() throws IOException, InterruptedException {
         if (AWS_ACCESS_KEY_ID == null) {
             return;
         }
+        // enable agentic search
+        updateClusterSettings(ML_COMMONS_AGENTIC_SEARCH_ENABLED.getKey(), true);
         queryPlanningModelId = registerQueryPlanningModel();
     }
 
@@ -92,8 +94,6 @@ public void testAgentWithQueryPlanningTool_DefaultPrompt() throws IOException {
         assertNotNull(agentId);
 
         String query = "{\"parameters\": {\"query_text\": \"How many iris flowers of type setosa are there?\"}}";
-        // enable agentic search
-        updateClusterSettings(ML_COMMONS_AGENTIC_SEARCH_ENABLED.getKey(), true);
         Response response = executeAgent(agentId, query);
         String responseBody = TestHelper.httpEntityToString(response.getEntity());
 

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,8 @@ public void setup() throws IOException, InterruptedException {`
`74`	`74`	`if (AWS_ACCESS_KEY_ID == null) {`
`75`	`75`	`return;`
`76`	`76`	`}`
	`77`	`+ // enable agentic search`
	`78`	`+ updateClusterSettings(ML_COMMONS_AGENTIC_SEARCH_ENABLED.getKey(), true);`
`77`	`79`	`queryPlanningModelId = registerQueryPlanningModel();`
`78`	`80`	`}`
`79`	`81`
`@@ -92,8 +94,6 @@ public void testAgentWithQueryPlanningTool_DefaultPrompt() throws IOException {`
`92`	`94`	`assertNotNull(agentId);`
`93`	`95`
`94`	`96`	`String query = "{\"parameters\": {\"query_text\": \"How many iris flowers of type setosa are there?\"}}";`
`95`		`- // enable agentic search`
`96`		`- updateClusterSettings(ML_COMMONS_AGENTIC_SEARCH_ENABLED.getKey(), true);`
`97`	`97`	`Response response = executeAgent(agentId, query);`
`98`	`98`	`String responseBody = TestHelper.httpEntityToString(response.getEntity());`
`99`	`99`