diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 312296c37..c02772644 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -26,6 +26,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD; import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW; @@ -183,14 +184,18 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); Map responseMap = entityAsMap(response); String workflowId = (String) responseMap.get(WORKFLOW_ID); - getAndAssertWorkflowStatus(workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); + // wait and ensure state is completed/done + assertBusy( + () -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + 30, + TimeUnit.SECONDS + ); // Hit Search State API with the workflow id created above String query = "{\"query\":{\"ids\":{\"values\":[\"" + workflowId + "\"]}}}"; SearchResponse searchResponse = searchWorkflowState(query); assertEquals(1, searchResponse.getHits().getTotalHits().value); String searchHitSource = searchResponse.getHits().getAt(0).getSourceAsString(); - System.out.println("search string: " + searchHitSource); WorkflowState searchHitWorkflowState = WorkflowState.parse(searchHitSource); // Assert based on the agent-framework template