From 9078d04b8d475b43dc12d30ce9c9361d3349ce01 Mon Sep 17 00:00:00 2001 From: Amit Galitzky Date: Fri, 15 Mar 2024 14:15:20 -0700 Subject: [PATCH] adding IT and UT Signed-off-by: Amit Galitzky --- .../flowframework/common/DefaultUseCases.java | 6 --- .../FlowFrameworkRestTestCase.java | 18 ++++++++ .../common/DefaultUseCasesTests.java | 44 +++++++++++++++++++ .../rest/FlowFrameworkRestApiIT.java | 41 +++++++++++++++++ 4 files changed, 103 insertions(+), 6 deletions(-) create mode 100644 src/test/java/org/opensearch/flowframework/common/DefaultUseCasesTests.java diff --git a/src/main/java/org/opensearch/flowframework/common/DefaultUseCases.java b/src/main/java/org/opensearch/flowframework/common/DefaultUseCases.java index baf23d884..1279ab209 100644 --- a/src/main/java/org/opensearch/flowframework/common/DefaultUseCases.java +++ b/src/main/java/org/opensearch/flowframework/common/DefaultUseCases.java @@ -38,12 +38,6 @@ public enum DefaultUseCases { "local_neural_sparse_search", "defaults/local-sparse-search-defaults.json", "substitutionTemplates/neural-sparse-local-template.json" - ), - /** defaults file and substitution ready template for cohere embedding model */ // TODO: not finalized - COHERE_EMBEDDING_MODEL_DEPLOY_SEMANTIC_SEARCH( - "cohere-embedding_model_deploy_semantic_search", - "defaults/cohere-embedding-defaults.json", - "substitutionTemplates/deploy-remote-model-template-v1.json" ); private final String useCaseName; diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index 37eeb14c7..326d382ee 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java @@ -339,6 +339,24 @@ protected Response createWorkflow(RestClient client, Template template) throws E return TestHelpers.makeRequest(client, "POST", WORKFLOW_URI + "?validation=off", Collections.emptyMap(), template.toJson(), null); } + /** + * Helper method to invoke the Create Workflow Rest Action without validation + * @param client the rest client + * @param useCase the usecase to create + * @throws Exception if the request fails + * @return a rest response + */ + protected Response createWorkflowWithUseCase(RestClient client, String useCase) throws Exception { + return TestHelpers.makeRequest( + client, + "POST", + WORKFLOW_URI + "?validation=off&use_case=" + useCase, + Collections.emptyMap(), + "{}", + null + ); + } + /** * Helper method to invoke the Create Workflow Rest Action with provision * @param client the rest client diff --git a/src/test/java/org/opensearch/flowframework/common/DefaultUseCasesTests.java b/src/test/java/org/opensearch/flowframework/common/DefaultUseCasesTests.java new file mode 100644 index 000000000..b6dc72ebb --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/common/DefaultUseCasesTests.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.common; + +import org.opensearch.flowframework.exception.FlowFrameworkException; +import org.opensearch.test.OpenSearchTestCase; + +public class DefaultUseCasesTests extends OpenSearchTestCase { + + @Override + public void setUp() throws Exception { + super.setUp(); + } + + public void testGetDefaultsFileByValidUseCaseName() throws FlowFrameworkException { + String defaultsFile = DefaultUseCases.getDefaultsFileByUseCaseName("open_ai_embedding_model_deploy"); + assertEquals("defaults/open-ai-embedding-defaults.json", defaultsFile); + } + + public void testGetDefaultsFileByInvalidUseCaseName() throws FlowFrameworkException { + FlowFrameworkException e = assertThrows( + FlowFrameworkException.class, + () -> DefaultUseCases.getDefaultsFileByUseCaseName("invalid_use_case") + ); + } + + public void testGetSubstitutionTemplateByValidUseCaseName() throws FlowFrameworkException { + String templateFile = DefaultUseCases.getSubstitutionReadyFileByUseCaseName("open_ai_embedding_model_deploy"); + assertEquals("substitutionTemplates/deploy-remote-model-template.json", templateFile); + } + + public void testGetSubstitutionTemplateByInvalidUseCaseName() throws FlowFrameworkException { + FlowFrameworkException e = assertThrows( + FlowFrameworkException.class, + () -> DefaultUseCases.getSubstitutionReadyFileByUseCaseName("invalid_use_case") + ); + } +} diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 69391cbd1..8db37d83d 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -399,4 +399,45 @@ public void testCreateAndProvisionIngestAndSearchPipeline() throws Exception { } + public void testDefaultCohereUseCase() throws Exception { + + // Using a 3 step template to create a connector, register remote model and deploy model + Template template = TestHelpers.createTemplateFromFile("ingest-search-pipeline-template.json"); + + // Hit Create Workflow API with original template + Response response = createWorkflowWithUseCase(client(), "cohere-embedding_model_deploy"); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + + Map responseMap = entityAsMap(response); + 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); + } + + assertEquals(RestStatus.OK, TestHelpers.restStatus(response)); + getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS); + + // Wait until provisioning has completed successfully before attempting to retrieve created resources + List resourcesCreated = getResourcesCreated(client(), workflowId, 30); + + List expectedStepNames = List.of("create_connector", "register_remote_model", "deploy_model"); + + List workflowStepNames = resourcesCreated.stream() + .peek(resourceCreated -> assertNotNull(resourceCreated.resourceId())) + .map(ResourceCreated::workflowStepName) + .collect(Collectors.toList()); + for (String expectedName : expectedStepNames) { + assertTrue(workflowStepNames.contains(expectedName)); + } + + // This template should create 5 resources, connector_id, registered model_id, deployed model_id and pipelineId + assertEquals(3, resourcesCreated.size()); + } + }