add more timeout

Signed-off-by: Jackie Han <[email protected]>
jackiehanyang · Dec 22, 2023 · 35c9ce1 · 35c9ce1
1 parent bea24d2
commit 35c9ce1
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
@@ -394,7 +394,7 @@ protected Response deprovisionWorkflow(String workflowId) throws Exception {
     protected Response deleteWorkflow(String workflowId) throws Exception {
         return TestHelpers.makeRequest(
             client(),
-            "POST",
+            "DELETE",
             String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, workflowId),
             ImmutableMap.of(),
             "",

diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -24,9 +24,12 @@
 import org.opensearch.flowframework.model.WorkflowState;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW;
@@ -196,22 +199,34 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception {
 
         // Assert based on the agent-framework template
         List<ResourceCreated> resourcesCreated = searchHitWorkflowState.resourcesCreated();
+        Set<String> expectedStepNames = new HashSet<>();
+        expectedStepNames.add("root_agent");
+        expectedStepNames.add("sub_agent");
+        expectedStepNames.add("openAI_connector");
+        expectedStepNames.add("gpt-3.5-model");
+        expectedStepNames.add("deployed-gpt-3.5-model");
+        Set<String> stepNames = resourcesCreated.stream().map(ResourceCreated::workflowStepId).collect(Collectors.toSet());
+
         assertEquals(5, resourcesCreated.size());
-        assertEquals("register_agent", resourcesCreated.get(0).workflowStepName());
+        assertEquals(stepNames, expectedStepNames);
         assertNotNull(resourcesCreated.get(0).resourceId());
 
         // Hit Deprovision API
         Response deprovisionResponse = deprovisionWorkflow(workflowId);
         assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
-        getAndAssertWorkflowStatus(workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED);
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            30,
+            TimeUnit.SECONDS
+        );
 
         // Hit Delete API
         Response deleteResponse = deleteWorkflow(workflowId);
         assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
 
         // Search this workflow id in global_context index to make sure it's deleted
         SearchResponse searchResponseAfterDeletion = searchWorkflows(query);
-        assertEquals(0, searchResponse.getHits().getTotalHits().value);
+        assertBusy(() -> assertEquals(0, searchResponseAfterDeletion.getHits().getTotalHits().value), 30, TimeUnit.SECONDS);
     }
 
 }