Skip to content

Commit

Permalink
Fixing forbiddenAPI check, addressing PR comments
Browse files Browse the repository at this point in the history 
Signed-off-by: Joshua Palis <[email protected]>
  • Loading branch information
@joshpalis
joshpalis committed Aug 14, 2024
1 parent 0e87cd9 commit 019f45d
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 92 deletions.
3 changes: 1 addition & 2 deletions src/test/java/org/opensearch/flowframework/TestHelpers.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import org.apache.hc.core5.http.Header;
import org.apache.hc.core5.http.HttpEntity;
import org.apache.hc.core5.http.io.entity.StringEntity;
import org.apache.logging.log4j.util.Strings;
import org.opensearch.client.Request;
import org.opensearch.client.RequestOptions;
import org.opensearch.client.Response;
Expand Down Expand Up @@ -74,7 +73,7 @@ public static Response makeRequest(
String jsonEntity,
List<Header> headers
) throws IOException {
HttpEntity httpEntity = Strings.isBlank(jsonEntity) ? null : new StringEntity(jsonEntity, APPLICATION_JSON);
HttpEntity httpEntity = jsonEntity.isBlank() ? null : new StringEntity(jsonEntity, APPLICATION_JSON);
return makeRequest(client, method, endpoint, params, httpEntity, headers);
}

Expand Down
103 changes: 13 additions & 90 deletions src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,7 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

import static org.opensearch.flowframework.common.CommonValue.CREATE_CONNECTOR_CREDENTIAL_KEY;
Expand All @@ -48,15 +46,12 @@
import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;

public class FlowFrameworkRestApiIT extends FlowFrameworkRestTestCase {
private static AtomicBoolean waitToStart = new AtomicBoolean(true);

@Before
public void waitToStart() throws Exception {
// ML Commons cron job runs every 10 seconds and takes 20+ seconds to initialize .plugins-ml-config index
// Delay on the first attempt for 25 seconds to allow this initialization and prevent flaky tests
if (waitToStart.getAndSet(false)) {
CountDownLatch latch = new CountDownLatch(1);
latch.await(25, TimeUnit.SECONDS);
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
}
}

Expand Down Expand Up @@ -93,14 +88,7 @@ public void testFailedUpdateWorkflow() throws Exception {
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
Response provisionResponse;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
provisionResponse = provisionWorkflow(client(), workflowId);
} else {
provisionResponse = provisionWorkflow(client(), workflowId);
}
Response provisionResponse = provisionResponse = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(provisionResponse));
getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);

Expand All @@ -122,14 +110,7 @@ public void testUpdateWorkflowUsingFields() throws Exception {
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
Response provisionResponse;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
provisionResponse = provisionWorkflow(client(), workflowId);
} else {
provisionResponse = provisionWorkflow(client(), workflowId);
}
Response provisionResponse = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(provisionResponse));
getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);

Expand Down Expand Up @@ -259,14 +240,7 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception {
String workflowId = (String) responseMap.get(WORKFLOW_ID);
getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = provisionWorkflow(client(), workflowId);
} else {
response = provisionWorkflow(client(), workflowId);
}

response = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);

Expand Down Expand Up @@ -294,13 +268,7 @@ 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;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = createWorkflowWithProvision(client(), template);
} else {
response = createWorkflowWithProvision(client(), template);
}
Response response = createWorkflowWithProvision(client(), template);
assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);
Expand Down Expand Up @@ -368,13 +336,7 @@ public void testReprovisionWorkflow() throws Exception {
Template template = TestHelpers.createTemplateFromFile("registerremotemodel.json");

// Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter
Response response;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = createWorkflowWithProvision(client(), template);
} else {
response = createWorkflowWithProvision(client(), template);
}
Response response = createWorkflowWithProvision(client(), template);
assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);
Expand Down Expand Up @@ -473,13 +435,7 @@ public void testReprovisionWorkflowMidNodeAddition() throws Exception {
Template template = TestHelpers.createTemplateFromFile("registerremotemodel-createindex.json");

// Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter
Response response;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = createWorkflowWithProvision(client(), template);
} else {
response = createWorkflowWithProvision(client(), template);
}
Response response = createWorkflowWithProvision(client(), template);
assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);
Expand Down Expand Up @@ -551,13 +507,7 @@ public void testReprovisionWorkflowMidNodeAddition() throws Exception {
public void testReprovisionWithNoChange() throws Exception {
Template template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline-createindex.json");

Response response;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = createWorkflowWithProvision(client(), template);
} else {
response = createWorkflowWithProvision(client(), template);
}
Response response = createWorkflowWithProvision(client(), template);
assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);
Expand Down Expand Up @@ -594,13 +544,7 @@ public void testReprovisionWithNoChange() throws Exception {
public void testReprovisionWithDeletion() throws Exception {
Template template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline-createindex.json");

Response response;
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = createWorkflowWithProvision(client(), template);
} else {
response = createWorkflowWithProvision(client(), template);
}
Response response = createWorkflowWithProvision(client(), template);
assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
Map<String, Object> responseMap = entityAsMap(response);
String workflowId = (String) responseMap.get(WORKFLOW_ID);
Expand Down Expand Up @@ -696,14 +640,7 @@ public void testCreateAndProvisionIngestAndSearchPipeline() throws Exception {
String workflowId = (String) responseMap.get(WORKFLOW_ID);
getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = provisionWorkflow(client(), workflowId);
} else {
response = provisionWorkflow(client(), workflowId);
}

response = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);

Expand Down Expand Up @@ -750,14 +687,7 @@ public void testDefaultCohereUseCase() throws Exception {
String workflowId = (String) responseMap.get(WORKFLOW_ID);
getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = provisionWorkflow(client(), workflowId);
} else {
response = provisionWorkflow(client(), workflowId);
}

response = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);

Expand Down Expand Up @@ -801,14 +731,7 @@ public void testDefaultSemanticSearchUseCaseWithFailureExpected() throws Excepti
String workflowId = (String) responseMap.get(WORKFLOW_ID);
getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED);

// Ensure Ml config index is initialized as creating a connector requires this, then hit Provision API and assert status
if (!indexExistsWithAdminClient(".plugins-ml-config")) {
assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS);
response = provisionWorkflow(client(), workflowId);
} else {
response = provisionWorkflow(client(), workflowId);
}

response = provisionWorkflow(client(), workflowId);
assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
getAndAssertWorkflowStatus(client(), workflowId, State.FAILED, ProvisioningProgress.FAILED);
}
Expand Down

0 comments on commit 019f45d

Please sign in to comment.