Fixing workflow data

Signed-off-by: Joshua Palis <jpalis@amazon.com>
opensearch-project · Sep 25, 2023 · 9662632 · 9662632
1 parent 8236580
commit 9662632
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 39 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIngestPipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIngestPipelineStep.java
@@ -125,12 +125,7 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
                 logger.info("Created ingest pipeline : " + putPipelineRequest.getId());
 
                 // PutPipelineRequest returns only an AcknowledgeResponse, returning pipelineId instead
-                createIngestPipelineFuture.complete(new WorkflowData() {
-                    @Override
-                    public Map<String, Object> getContent() {
-                        return Map.of("pipelineId", putPipelineRequest.getId());
-                    }
-                });
+                createIngestPipelineFuture.complete(new WorkflowData(Map.of("pipelineId", putPipelineRequest.getId())));
 
                 // TODO : Use node client to index response data to global context (pending global context index implementation)
 

diff --git a/...opensearch/flowframework/workflow/CreateIngestPipeline/CreateIngestPipelineStepTests.java b/...opensearch/flowframework/workflow/CreateIngestPipeline/CreateIngestPipelineStepTests.java
@@ -43,29 +43,19 @@ public class CreateIngestPipelineStepTests extends OpenSearchTestCase {
     public void setUp() throws Exception {
         super.setUp();
 
-        inputData = new WorkflowData() {
-
-            @Override
-            public Map<String, Object> getContent() {
-                return Map.ofEntries(
-                    Map.entry("id", "pipelineId"),
-                    Map.entry("description", "some description"),
-                    Map.entry("type", "text_embedding"),
-                    Map.entry("model_id", "model_id"),
-                    Map.entry("input_field_name", "inputField"),
-                    Map.entry("output_field_name", "outputField")
-                );
-            }
-        };
+        inputData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("id", "pipelineId"),
+                Map.entry("description", "some description"),
+                Map.entry("type", "text_embedding"),
+                Map.entry("model_id", "model_id"),
+                Map.entry("input_field_name", "inputField"),
+                Map.entry("output_field_name", "outputField")
+            )
+        );
 
         // Set output data to returned pipelineId
-        outpuData = new WorkflowData() {
-
-            @Override
-            public Map<String, Object> getContent() {
-                return Map.ofEntries(Map.entry("pipelineId", "pipelineId"));
-            }
-        };
+        outpuData = new WorkflowData(Map.ofEntries(Map.entry("pipelineId", "pipelineId")));
 
         client = mock(Client.class);
         adminClient = mock(AdminClient.class);
@@ -116,18 +106,14 @@ public void testMissingData() throws InterruptedException {
         CreateIngestPipelineStep createIngestPipelineStep = new CreateIngestPipelineStep(client);
 
         // Data with missing input and output fields
-        WorkflowData incorrectData = new WorkflowData() {
-
-            @Override
-            public Map<String, Object> getContent() {
-                return Map.ofEntries(
-                    Map.entry("id", "pipelineId"),
-                    Map.entry("description", "some description"),
-                    Map.entry("type", "text_embedding"),
-                    Map.entry("model_id", "model_id")
-                );
-            }
-        };
+        WorkflowData incorrectData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("id", "pipelineId"),
+                Map.entry("description", "some description"),
+                Map.entry("type", "text_embedding"),
+                Map.entry("model_id", "model_id")
+            )
+        );
 
         CompletableFuture<WorkflowData> future = createIngestPipelineStep.execute(List.of(incorrectData));
         assertTrue(future.isDone() && future.isCompletedExceptionally());