Adding reprovision integration tests

src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java

@@ -415,6 +415,25 @@ protected Response createWorkflowValidation(RestClient client, Template template
         return TestHelpers.makeRequest(client, "POST", WORKFLOW_URI, Collections.emptyMap(), template.toJson(), null);
     }
 
+    /**
+     * Helper method to invoke the Reprovision Workflow API
+     * @param client the rest client
+     * @param workflowId the document id
+     * @param templateFields the template to reprovision
+     * @throws Exception if the request fails
+     * @return a rest response
+     */
+    protected Response reprovisionWorkflow(RestClient client, String workflowId, Template template) throws Exception {
+        return TestHelpers.makeRequest(
+            client,
+            "PUT",
+            String.format(Locale.ROOT, "%s/%s?reprovision=true", WORKFLOW_URI, workflowId),
+            Collections.emptyMap(),
+            template.toJson(),
+            null
+        );
+    }
+
     /**
      * Helper method to invoke the Update Workflow API
      * @param client the rest client

src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java

@@ -363,6 +363,281 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception {
         assertBusy(() -> { getAndAssertWorkflowStatusNotFound(client(), workflowId); }, 30, TimeUnit.SECONDS);
     }
 
+    public void testReprovisionWorkflow() throws Exception {
+        // Begin with a template to register a local pretrained model
+        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);
+        }
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+        // wait and ensure state is completed/done
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); },
+            120,
+            TimeUnit.SECONDS
+        );
+
+        // Wait until provisioning has completed successfully before attempting to retrieve created resources
+        List<ResourceCreated> resourcesCreated = getResourcesCreated(client(), workflowId, 30);
+        assertEquals(3, resourcesCreated.size());
+        List<String> resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList());
+        assertTrue(resourceIds.contains("create_connector"));
+        assertTrue(resourceIds.contains("register_remote_model"));
+
+        // Reprovision template to add ingest pipeline which uses the model ID
+        template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline.json");
+        response = reprovisionWorkflow(client(), workflowId, template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
+        resourcesCreated = getResourcesCreated(client(), workflowId, 10);
+        assertEquals(4, resourcesCreated.size());
+        resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList());
+        assertTrue(resourceIds.contains("create_connector"));
+        assertTrue(resourceIds.contains("register_remote_model"));
+        assertTrue(resourceIds.contains("create_ingest_pipeline"));
+
+        // Retrieve pipeline by ID to ensure model ID is set correctly
+        String modelId = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("register_remote_model"))
+            .findFirst()
+            .get()
+            .resourceId();
+        String pipelineId = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("create_ingest_pipeline"))
+            .findFirst()
+            .get()
+            .resourceId();
+        GetPipelineResponse getPipelineResponse = getPipelines(pipelineId);
+        assertEquals(1, getPipelineResponse.pipelines().size());
+        assertTrue(getPipelineResponse.pipelines().get(0).getConfigAsMap().toString().contains(modelId));
+
+        // Reprovision template to add index which uses default ingest pipeline
+        template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline-createindex.json");
+        response = reprovisionWorkflow(client(), workflowId, template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
+        resourcesCreated = getResourcesCreated(client(), workflowId, 10);
+        assertEquals(5, resourcesCreated.size());
+        resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList());
+        assertTrue(resourceIds.contains("create_connector"));
+        assertTrue(resourceIds.contains("register_remote_model"));
+        assertTrue(resourceIds.contains("create_ingest_pipeline"));
+        assertTrue(resourceIds.contains("create_index"));
+
+        // Retrieve index settings to ensure pipeline ID is set correctly
+        String indexName = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("create_index"))
+            .findFirst()
+            .get()
+            .resourceId();
+        Map<String, Object> indexSettings = getIndexSettingsAsMap(indexName);
+        assertEquals(pipelineId, indexSettings.get("index.default_pipeline"));
+
+        // Reprovision template to remove default ingest pipeline
+        template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline-updateindex.json");
+        response = reprovisionWorkflow(client(), workflowId, template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
+        resourcesCreated = getResourcesCreated(client(), workflowId, 10);
+        // resource count should remain unchanged when updating an existing node
+        assertEquals(5, resourcesCreated.size());
+
+        // Retrieve index settings to ensure default pipeline has been updated correctly
+        indexSettings = getIndexSettingsAsMap(indexName);
+        assertEquals("_none", indexSettings.get("index.default_pipeline"));
+
+        // Deprovision and delete all resources
+        Response deprovisionResponse = deprovisionWorkflowWithAllowDelete(client(), workflowId, pipelineId + "," + indexName);
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            60,
+            TimeUnit.SECONDS
+        );
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
+
+        // Hit Delete API
+        Response deleteResponse = deleteWorkflow(client(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
+    }
+
+    public void testReprovisionWorkflowMidNodeAddition() throws Exception {
+        // Begin with a template to register a local pretrained model and create an index, no edges
+        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);
+        }
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+        // wait and ensure state is completed/done
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); },
+            120,
+            TimeUnit.SECONDS
+        );
+
+        // Wait until provisioning has completed successfully before attempting to retrieve created resources
+        List<ResourceCreated> resourcesCreated = getResourcesCreated(client(), workflowId, 30);
+        assertEquals(4, resourcesCreated.size());
+        List<String> resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList());
+        assertTrue(resourceIds.contains("create_connector"));
+        assertTrue(resourceIds.contains("register_remote_model"));
+        assertTrue(resourceIds.contains("create_index"));
+
+        // Reprovision template to add ingest pipeline which uses the model ID
+        template = TestHelpers.createTemplateFromFile("registerremotemodel-ingestpipeline-createindex.json");
+        response = reprovisionWorkflow(client(), workflowId, template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
+        resourcesCreated = getResourcesCreated(client(), workflowId, 10);
+        assertEquals(5, resourcesCreated.size());
+        resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList());
+        assertTrue(resourceIds.contains("create_connector"));
+        assertTrue(resourceIds.contains("register_remote_model"));
+        assertTrue(resourceIds.contains("create_ingest_pipeline"));
+        assertTrue(resourceIds.contains("create_index"));
+
+        // Ensure ingest pipeline configuration contains the model id and index settings have the ingest pipeline as default
+        String modelId = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("register_remote_model"))
+            .findFirst()
+            .get()
+            .resourceId();
+        String pipelineId = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("create_ingest_pipeline"))
+            .findFirst()
+            .get()
+            .resourceId();
+        GetPipelineResponse getPipelineResponse = getPipelines(pipelineId);
+        assertEquals(1, getPipelineResponse.pipelines().size());
+        assertTrue(getPipelineResponse.pipelines().get(0).getConfigAsMap().toString().contains(modelId));
+
+        String indexName = resourcesCreated.stream()
+            .filter(x -> x.workflowStepName().equals("create_index"))
+            .findFirst()
+            .get()
+            .resourceId();
+        Map<String, Object> indexSettings = getIndexSettingsAsMap(indexName);
+        assertEquals(pipelineId, indexSettings.get("index.default_pipeline"));
+
+        // Deprovision and delete all resources
+        Response deprovisionResponse = deprovisionWorkflowWithAllowDelete(client(), workflowId, pipelineId + "," + indexName);
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            60,
+            TimeUnit.SECONDS
+        );
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
+
+        // Hit Delete API
+        Response deleteResponse = deleteWorkflow(client(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
+    }
+
+    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);
+        }
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+        // wait and ensure state is completed/done
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); },
+            120,
+            TimeUnit.SECONDS
+        );
+
+        // Attempt to reprovision the same template with no changes
+        ResponseException exception = expectThrows(ResponseException.class, () -> reprovisionWorkflow(client(), workflowId, template));
+        assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode());
+        assertTrue(exception.getMessage().contains("Template does not contain any modifications"));
+
+        // Deprovision and delete all resources
+        Response deprovisionResponse = deprovisionWorkflowWithAllowDelete(
+            client(),
+            workflowId,
+            "nlp-ingest-pipeline" + "," + "my-nlp-index"
+        );
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            60,
+            TimeUnit.SECONDS
+        );
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
+
+        // Hit Delete API
+        Response deleteResponse = deleteWorkflow(client(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
+    }
+
+    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);
+        }
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+        // wait and ensure state is completed/done
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); },
+            120,
+            TimeUnit.SECONDS
+        );
+
+        // Attempt to reprovision template without ingest pipeline node
+        Template templateWithoutIngestPipeline = TestHelpers.createTemplateFromFile("registerremotemodel-createindex.json");
+        ResponseException exception = expectThrows(
+            ResponseException.class,
+            () -> reprovisionWorkflow(client(), workflowId, templateWithoutIngestPipeline)
+        );
+        assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode());
+        assertTrue(exception.getMessage().contains("Workflow Step deletion is not supported when reprovisioning a template."));
+
+        // Deprovision and delete all resources
+        Response deprovisionResponse = deprovisionWorkflowWithAllowDelete(
+            client(),
+            workflowId,
+            "nlp-ingest-pipeline" + "," + "my-nlp-index"
+        );
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(client(), workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            60,
+            TimeUnit.SECONDS
+        );
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
+
+        // Hit Delete API
+        Response deleteResponse = deleteWorkflow(client(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
+    }
+
     public void testTimestamps() throws Exception {
         Template noopTemplate = TestHelpers.createTemplateFromFile("noop.json");
         // Create the template, should have created and updated matching

src/test/resources/template/registerremotemodel-createindex.json

@@ -0,0 +1,84 @@
+{
+    "name": "semantic search with local pretrained model",
+    "description": "Setting up semantic search, with a local pretrained embedding model",
+    "use_case": "SEMANTIC_SEARCH",
+    "version": {
+      "template": "1.0.0",
+      "compatibility": [
+        "2.12.0",
+        "3.0.0"
+      ]
+    },
+    "workflows": {
+      "provision": {
+        "nodes": [
+          {
+            "id": "create_openai_connector",
+            "type": "create_connector",
+            "user_inputs": {
+              "name": "OpenAI Chat Connector",
+              "description": "The connector to public OpenAI model service for text embedding model",
+              "version": "1",
+              "protocol": "http",
+              "parameters": {
+                "endpoint": "api.openai.com",
+                "model": "gpt-3.5-turbo",
+                "response_filter": "$.choices[0].message.content"
+              },
+              "credential": {
+                "openAI_key": "12345"
+              },
+              "actions": [
+                {
+                  "action_type": "predict",
+                  "method": "POST",
+                  "url": "https://${parameters.endpoint}/v1/chat/completions"
+                }
+              ]
+            }
+          },
+          {
+            "id": "register_openai_model",
+            "type": "register_remote_model",
+            "previous_node_inputs": {
+              "create_openai_connector": "connector_id"
+            },
+            "user_inputs": {
+              "name": "openAI-gpt-3.5-turbo",
+              "deploy": true
+            }
+          },
+          {
+            "id": "create_index",
+            "type": "create_index",
+            "user_inputs": {
+              "index_name": "my-nlp-index",
+              "configurations": {
+                "settings": {
+                  "index.knn": true,
+                  "index.number_of_shards": "2"
+                },
+                "mappings": {
+                  "properties": {
+                    "passage_embedding": {
+                      "type": "knn_vector",
+                      "dimension": "768",
+                      "method": {
+                        "engine": "lucene",
+                        "space_type": "l2",
+                        "name": "hnsw",
+                        "parameters": {}
+                      }
+                    },
+                    "passage_text": {
+                      "type": "text"
+                    }
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    }
+  }