From 60a10b3deee49479d537729eedd3888ae0c7690f Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Thu, 21 Dec 2023 15:22:04 -0800 Subject: [PATCH 1/5] add integ tests for search state api, deprovision api, and delete api Signed-off-by: Jackie Han --- .../flowframework/model/WorkflowState.java | 19 ++ .../FlowFrameworkRestTestCase.java | 63 ++++++ .../rest/FlowFrameworkRestApiIT.java | 41 +++- .../resources/template/agent-framework.json | 201 ++++++++++++++++++ 4 files changed, 323 insertions(+), 1 deletion(-) create mode 100644 src/test/resources/template/agent-framework.json diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java index b8086aa4c..56f5d4e2d 100644 --- a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java +++ b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java @@ -8,12 +8,15 @@ */ package org.opensearch.flowframework.model; +import org.opensearch.common.xcontent.LoggingDeprecationHandler; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.commons.authuser.User; import org.opensearch.core.common.ParsingException; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParseException; @@ -393,6 +396,22 @@ public static WorkflowState parse(XContentParser parser) throws IOException { .build(); } + /** + * Parse a JSON workflow state + * @param json A string containing a JSON representation of a workflow state + * @return A {@link WorkflowState} represented by the JSON + * @throws IOException on failure to parse + */ + public static WorkflowState parse(String json) throws IOException { + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + json + ); + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); + return parse(parser); + } + /** * The workflowID associated with this workflow-state * @return the workflowId diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index ac537047f..f34fdc472 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java @@ -320,6 +320,10 @@ protected Response createWorkflow(Template template) throws Exception { return TestHelpers.makeRequest(client(), "POST", WORKFLOW_URI, ImmutableMap.of(), template.toJson(), null); } + protected Response createWorkflowWithProvision(Template template) throws Exception { + return TestHelpers.makeRequest(client(), "POST", WORKFLOW_URI + "?provision=true", ImmutableMap.of(), template.toJson(), null); + } + /** * Helper method to invoke the Create Workflow Rest Action with dry run validation * @param template the template to create @@ -365,6 +369,40 @@ protected Response provisionWorkflow(String workflowId) throws Exception { ); } + /** + * Helper method to invoke the Deprovision Workflow Rest Action + * @param workflowId the workflow ID to deprovision + * @return a rest response + * @throws Exception if the request fails + */ + protected Response deprovisionWorkflow(String workflowId) throws Exception { + return TestHelpers.makeRequest( + client(), + "POST", + String.format(Locale.ROOT, "%s/%s/%s", WORKFLOW_URI, workflowId, "_deprovision"), + ImmutableMap.of(), + "", + null + ); + } + + /** + * Helper method to invoke the Delete Workflow Rest Action + * @param workflowId the workflow ID to delete + * @return a rest response + * @throws Exception if the request fails + */ + protected Response deleteWorkflow(String workflowId) throws Exception { + return TestHelpers.makeRequest( + client(), + "POST", + String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, workflowId), + ImmutableMap.of(), + "", + null + ); + } + /** * Helper method to invoke the Get Workflow Rest Action * @param workflowId the workflow ID to get the status @@ -417,6 +455,31 @@ protected SearchResponse searchWorkflows(String query) throws Exception { } } + protected SearchResponse searchWorkflowState(String query) throws Exception { + Response restSearchResponse = TestHelpers.makeRequest( + client(), + "GET", + String.format(Locale.ROOT, "%s/state/_search", WORKFLOW_URI), + ImmutableMap.of(), + query, + null + ); + assertEquals(RestStatus.OK, TestHelpers.restStatus(restSearchResponse)); + + // Parse entity content into SearchResponse + MediaType mediaType = MediaType.fromMediaType(restSearchResponse.getEntity().getContentType()); + try ( + XContentParser parser = mediaType.xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + restSearchResponse.getEntity().getContent() + ) + ) { + return SearchResponse.fromXContent(parser); + } + } + /** * Helper method to invoke the Get Workflow Rest Action and assert the provisioning and state status * @param workflowId the workflow ID to get the status diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index e8a99946e..10d382727 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -21,6 +21,7 @@ import org.opensearch.flowframework.model.Workflow; import org.opensearch.flowframework.model.WorkflowEdge; import org.opensearch.flowframework.model.WorkflowNode; +import org.opensearch.flowframework.model.WorkflowState; import java.util.Collections; import java.util.HashMap; @@ -168,7 +169,7 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception { getAndAssertWorkflowStatus(workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); // Wait until provisioning has completed successfully before attempting to retrieve created resources - List resourcesCreated = getResourcesCreated(workflowId, 10); + List resourcesCreated = getResourcesCreated(workflowId, 20); // This template should create 3 resources, connector_id, regestered model_id and deployed model_id assertEquals(3, resourcesCreated.size()); @@ -176,4 +177,42 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception { assertNotNull(resourcesCreated.get(0).resourceId()); } + public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { + Template template = TestHelpers.createTemplateFromFile("agent-framework.json"); + + // Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter + Response response = createWorkflowWithProvision(template); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String workflowId = (String) responseMap.get(WORKFLOW_ID); + getAndAssertWorkflowStatus(workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); + + // 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 + List resourcesCreated = searchHitWorkflowState.resourcesCreated(); + assertEquals(5, resourcesCreated.size()); + assertEquals("register_agent", resourcesCreated.get(0).workflowStepName()); + 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); + + // 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); + } + } diff --git a/src/test/resources/template/agent-framework.json b/src/test/resources/template/agent-framework.json new file mode 100644 index 000000000..224db271c --- /dev/null +++ b/src/test/resources/template/agent-framework.json @@ -0,0 +1,201 @@ +{ + "name": "opensearch-assistant-workflow", + "description": "test case", + "use_case": "REGISTER_AGENT", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "openAI_connector", + "type": "create_connector", + "user_inputs": { + "name": "OpenAI Chat Connector", + "description": "The connector to public OpenAI model service for GPT 3.5", + "protocol": "aws_sigv4", + "version": "1", + "parameters": { + "region": "us-east-1", + "service_name": "bedrock", + "anthropic_version": "bedrock-2023-05-31", + "endpoint": "bedrock.us-east-1.amazonaws.com", + "auth": "Sig_V4", + "content_type": "application/json", + "max_tokens_to_sample": 8000, + "temperature": 0.0001 + }, + "credential": { + "access_key": "", + "secret_key": "", + "session_token": "" + }, + "actions": [ + { + "action_type": "predict", + "method": "POST", + "url": "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/invoke", + "request_body": "{\"prompt\":\"${parameters.prompt}\", \"max_tokens_to_sample\":${parameters.max_tokens_to_sample}, \"temperature\":${parameters.temperature}, \"anthropic_version\":\"${parameters.anthropic_version}\" }" + } + ] + } + }, + { + "id": "gpt-3.5-model", + "type": "register_remote_model", + "previous_node_inputs": { + "openAI_connector": "connector_id" + }, + "user_inputs": { + "name": "flow-register-remote-test-gpt-3.5-turbo", + "function_name": "remote", + "description": "test model" + } + }, + { + "id": "deployed-gpt-3.5-model", + "type": "deploy_model", + "previous_node_inputs": { + "gpt-3.5-model": "model_id" + } + }, + { + "id": "ml_model_tool", + "type": "create_tool", + "previous_node_inputs": { + "deployed-gpt-3.5-model": "model_id" + }, + "user_inputs": { + "name": "MLModelTool", + "type": "MLModelTool", + "alias": "language_model_tool", + "description": "A general tool to answer any question. But it can't handle math problem", + "parameters": { + "prompt": "Answer the question as best you can.", + "response_filter": "choices[0].message.content" + } + } + }, + { + "id": "math_tool", + "type": "create_tool", + "user_inputs": { + "name": "MathTool", + "type": "MathTool", + "description": "A general tool to calculate any math problem. The action input must be valid math expression, like 2+3", + "parameters": { + "max_iteration": 5 + } + } + }, + { + "id": "sub_agent", + "type": "register_agent", + "previous_node_inputs": { + "math_tool": "tools" + }, + "user_inputs": { + "name": "Sub Agent", + "type": "cot", + "description": "this is a test agent", + "parameters": { + "hello": "world" + }, + "llm.model_id": "ldzS04kBxRPZ5cnWrqpd", + "llm.parameters": { + "max_iteration": "5", + "stop_when_no_tool_found": "true" + }, + "memory": { + "type": "conversation_index" + }, + "app_type": "chatbot", + "created_time": 1689793598499, + "last_updated_time": 1689793598530 + } + }, + { + "id": "agent_tool", + "type": "create_tool", + "previous_node_inputs": { + "sub_agent": "agent_id" + }, + "user_inputs": { + "name": "AgentTool", + "type": "AgentTool", + "description": "Root Agent Tool", + "parameters": { + "max_iteration": 5 + } + } + }, + { + "id": "root_agent", + "type": "register_agent", + "previous_node_inputs": { + "deployed-gpt-3.5-model": "model_id", + "ml_model_tool": "tools", + "agent_tool": "tools" + }, + "user_inputs": { + "name": "DEMO-Test_Agent_For_CoT", + "type": "cot", + "description": "this is a test agent", + "parameters": { + "hello": "world" + }, + "llm.parameters": { + "max_iteration": "5", + "stop_when_no_tool_found": "true" + }, + "memory": { + "type": "conversation_index" + }, + "app_type": "chatbot", + "created_time": 1689793598499, + "last_updated_time": 1689793598530 + } + } + ], + "edges": [ + { + "source": "openAI_connector", + "dest": "gpt-3.5-model" + }, + { + "source": "gpt-3.5-model", + "dest": "deployed-gpt-3.5-model" + }, + { + "source": "deployed-gpt-3.5-model", + "dest": "root_agent" + }, + { + "source": "deployed-gpt-3.5-model", + "dest": "ml_model_tool" + }, + { + "source": "ml_model_tool", + "dest": "root_agent" + }, + { + "source": "math_tool", + "dest": "sub_agent" + }, + { + "source": "sub_agent", + "dest": "agent_tool" + }, + { + "source": "agent_tool", + "dest": "root_agent" + } + ] + } + } +} \ No newline at end of file From a34dcd3524db1917a2a4bc4629d1475e11d3e7f1 Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Fri, 22 Dec 2023 11:12:45 -0800 Subject: [PATCH 2/5] spotless check clean up Signed-off-by: Jackie Han --- .../flowframework/model/WorkflowState.java | 6 +-- .../FlowFrameworkRestTestCase.java | 48 +++++++++---------- .../resources/template/agent-framework.json | 2 +- 3 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java index 56f5d4e2d..e962f5758 100644 --- a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java +++ b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java @@ -404,9 +404,9 @@ public static WorkflowState parse(XContentParser parser) throws IOException { */ public static WorkflowState parse(String json) throws IOException { XContentParser parser = JsonXContent.jsonXContent.createParser( - NamedXContentRegistry.EMPTY, - LoggingDeprecationHandler.INSTANCE, - json + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + json ); ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); return parse(parser); diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index c4ff4eba5..f20f47e33 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java @@ -376,12 +376,12 @@ protected Response provisionWorkflow(String workflowId) throws Exception { */ protected Response deprovisionWorkflow(String workflowId) throws Exception { return TestHelpers.makeRequest( - client(), - "POST", - String.format(Locale.ROOT, "%s/%s/%s", WORKFLOW_URI, workflowId, "_deprovision"), - ImmutableMap.of(), - "", - null + client(), + "POST", + String.format(Locale.ROOT, "%s/%s/%s", WORKFLOW_URI, workflowId, "_deprovision"), + ImmutableMap.of(), + "", + null ); } @@ -393,12 +393,12 @@ protected Response deprovisionWorkflow(String workflowId) throws Exception { */ protected Response deleteWorkflow(String workflowId) throws Exception { return TestHelpers.makeRequest( - client(), - "POST", - String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, workflowId), - ImmutableMap.of(), - "", - null + client(), + "POST", + String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, workflowId), + ImmutableMap.of(), + "", + null ); } @@ -456,24 +456,24 @@ protected SearchResponse searchWorkflows(String query) throws Exception { protected SearchResponse searchWorkflowState(String query) throws Exception { Response restSearchResponse = TestHelpers.makeRequest( - client(), - "GET", - String.format(Locale.ROOT, "%s/state/_search", WORKFLOW_URI), - ImmutableMap.of(), - query, - null + client(), + "GET", + String.format(Locale.ROOT, "%s/state/_search", WORKFLOW_URI), + ImmutableMap.of(), + query, + null ); assertEquals(RestStatus.OK, TestHelpers.restStatus(restSearchResponse)); // Parse entity content into SearchResponse MediaType mediaType = MediaType.fromMediaType(restSearchResponse.getEntity().getContentType()); try ( - XContentParser parser = mediaType.xContent() - .createParser( - NamedXContentRegistry.EMPTY, - DeprecationHandler.THROW_UNSUPPORTED_OPERATION, - restSearchResponse.getEntity().getContent() - ) + XContentParser parser = mediaType.xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + restSearchResponse.getEntity().getContent() + ) ) { return SearchResponse.fromXContent(parser); } diff --git a/src/test/resources/template/agent-framework.json b/src/test/resources/template/agent-framework.json index 224db271c..713214638 100644 --- a/src/test/resources/template/agent-framework.json +++ b/src/test/resources/template/agent-framework.json @@ -198,4 +198,4 @@ ] } } -} \ No newline at end of file +} From bea24d2c864fc0d40cc20dacd394e365a828e410 Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Fri, 22 Dec 2023 11:57:50 -0800 Subject: [PATCH 3/5] add timeout Signed-off-by: Jackie Han --- .../flowframework/rest/FlowFrameworkRestApiIT.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 312296c37..e4e827569 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,14 @@ 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 From 35c9ce120f9e18dde36a9141de9695365389cadd Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Fri, 22 Dec 2023 12:32:02 -0800 Subject: [PATCH 4/5] add more timeout Signed-off-by: Jackie Han --- .../FlowFrameworkRestTestCase.java | 2 +- .../rest/FlowFrameworkRestApiIT.java | 21 ++++++++++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index f20f47e33..15ffb65e1 100644 --- 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 index e4e827569..6999d2779 100644 --- 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,14 +199,26 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { // Assert based on the agent-framework template List resourcesCreated = searchHitWorkflowState.resourcesCreated(); + Set 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 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); @@ -211,7 +226,7 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { // 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); } } From 258d4fa9370acd66c8a0635c00655525c7f09bf6 Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Fri, 22 Dec 2023 13:51:00 -0800 Subject: [PATCH 5/5] Add thread sleep to wait for deletion to complete Signed-off-by: Jackie Han --- .../opensearch/flowframework/rest/FlowFrameworkRestApiIT.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 6999d2779..cfc5919d4 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -224,9 +224,13 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { Response deleteResponse = deleteWorkflow(workflowId); assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse)); + // wait for deletion to complete + Thread.sleep(30000); + // Search this workflow id in global_context index to make sure it's deleted SearchResponse searchResponseAfterDeletion = searchWorkflows(query); assertBusy(() -> assertEquals(0, searchResponseAfterDeletion.getHits().getTotalHits().value), 30, TimeUnit.SECONDS); + } }