From 9e3abdc0fdcbc440962cfae8281b9f5aa77b509e Mon Sep 17 00:00:00 2001 From: Amit Galitzky Date: Fri, 12 Jan 2024 12:10:26 -0800 Subject: [PATCH] testing with waitUntil Signed-off-by: Amit Galitzky --- .../FlowFrameworkRestTestCase.java | 31 +++++++++++++++--- .../rest/FlowFrameworkRestApiIT.java | 32 ++++++++++--------- 2 files changed, 43 insertions(+), 20 deletions(-) diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index a0badfa15..659dbe7f2 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java @@ -61,6 +61,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import static org.opensearch.client.RestClientBuilder.DEFAULT_MAX_CONN_PER_ROUTE; import static org.opensearch.client.RestClientBuilder.DEFAULT_MAX_CONN_TOTAL; @@ -490,6 +491,25 @@ protected void getAndAssertWorkflowStatus(String workflowId, State stateStatus, } + public void waitForProvisioningStatus(String workflowId, ProvisioningProgress provisioningProgress) throws InterruptedException { + AtomicBoolean provisioningStatusDone = new AtomicBoolean(false); + waitUntil(() -> { + try { + Response response = getWorkflowStatus(workflowId, true); + assertEquals(RestStatus.OK, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String state = (String) responseMap.get(CommonValue.PROVISIONING_PROGRESS_FIELD); + if (provisioningProgress.name().equals(state)) { + provisioningStatusDone.set(true); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + return provisioningStatusDone.get(); + }, 90, TimeUnit.SECONDS); + assertTrue(provisioningStatusDone.get()); + } + protected void getAndAssertWorkflowStep() throws Exception { Response response = getWorkflowStep(); assertEquals(RestStatus.OK, TestHelpers.restStatus(response)); @@ -505,11 +525,12 @@ protected void getAndAssertWorkflowStep() throws Exception { protected List getResourcesCreated(String workflowId, int timeout) throws Exception { // wait and ensure state is completed/done - assertBusy( - () -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, - timeout, - TimeUnit.SECONDS - ); + // assertBusy( + // () -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + // timeout, + // TimeUnit.SECONDS + // ); + waitForProvisioningStatus(workflowId, ProvisioningProgress.DONE); Response response = getWorkflowStatus(workflowId, true); diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 75dd3a8f1..814b23751 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -29,7 +29,6 @@ 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; @@ -188,11 +187,14 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception { response = provisionWorkflow(workflowId); assertEquals(RestStatus.OK, TestHelpers.restStatus(response)); Thread.sleep(1000); - assertBusy( - () -> { getAndAssertWorkflowStatus(workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); }, - 30, - TimeUnit.SECONDS - ); + + // waitForProvisioningStatus(workflowId, ProvisioningProgress.DONE); + + // assertBusy( + // () -> { getAndAssertWorkflowStatus(workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); }, + // 30, + // TimeUnit.SECONDS + // ); // Wait until provisioning has completed successfully before attempting to retrieve created resources List resourcesCreated = getResourcesCreated(workflowId, 90); @@ -207,7 +209,7 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception { assertNotNull(resourcesCreated.get(2).resourceId()); // // Deprovision the workflow to avoid opening circut breaker when running additional tests - // Response deprovisionResponse = deprovisionWorkflow(workflowId); + Response deprovisionResponse = deprovisionWorkflow(workflowId); // wait for deprovision to complete Thread.sleep(5000); @@ -222,8 +224,8 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { Map responseMap = entityAsMap(response); String workflowId = (String) responseMap.get(WORKFLOW_ID); // wait and ensure state is completed/done - assertBusy(() -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, 60, TimeUnit.SECONDS); - + // assertBusy(() -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, 60, TimeUnit.SECONDS); + waitForProvisioningStatus(workflowId, ProvisioningProgress.DONE); // Hit Search State API with the workflow id created above String query = "{\"query\":{\"ids\":{\"values\":[\"" + workflowId + "\"]}}}"; SearchResponse searchResponse = searchWorkflowState(query); @@ -246,12 +248,12 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { // Hit Deprovision API Response deprovisionResponse = deprovisionWorkflow(workflowId); - assertBusy( - () -> { getAndAssertWorkflowStatus(workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); }, - 60, - TimeUnit.SECONDS - ); - + // assertBusy( + // () -> { getAndAssertWorkflowStatus(workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); }, + // 60, + // TimeUnit.SECONDS + // ); + waitForProvisioningStatus(workflowId, ProvisioningProgress.NOT_STARTED); // Hit Delete API Response deleteResponse = deleteWorkflow(workflowId); assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));