From 99b2412ad4ad3af72bff189404f6d817417660b8 Mon Sep 17 00:00:00 2001
From: Chris Sampson <chris.sampson82@gmail.com>
Date: Tue, 25 Apr 2023 21:18:02 +0100
Subject: [PATCH] NIFI-11483 Correctly use DescribedValue for
 JsonQueryElasticsearch

This closes #7195

Signed-off-by: David Handermann <exceptionfactory@apache.org>
---
 .../elasticsearch/AbstractJsonQueryElasticsearch.java |  4 ++--
 .../elasticsearch/api/ResultOutputStrategy.java       |  5 +++++
 .../AbstractJsonQueryElasticsearchTest.groovy         | 11 ++++++++---
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearch.java b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearch.java
index a0d32779b0..1000779143 100644
--- a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearch.java
+++ b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearch.java
@@ -65,7 +65,7 @@ public abstract class AbstractJsonQueryElasticsearch<Q extends JsonQueryParamete
             .name("el-rest-split-up-hits")
             .displayName("Search Results Split")
             .description("Output a flowfile containing all hits or one flowfile for each individual hit.")
-            .allowableValues(ResultOutputStrategy.PER_RESPONSE.getValue(), ResultOutputStrategy.PER_HIT.getValue())
+            .allowableValues(ResultOutputStrategy.getNonPaginatedResponseOutputStrategies())
             .defaultValue(ResultOutputStrategy.PER_RESPONSE.getValue())
             .required(true)
             .expressionLanguageSupported(ExpressionLanguageScope.NONE)
@@ -84,7 +84,7 @@ public abstract class AbstractJsonQueryElasticsearch<Q extends JsonQueryParamete
             .name("el-rest-split-up-aggregations")
             .displayName("Aggregation Results Split")
             .description("Output a flowfile containing all aggregations or one flowfile for each individual aggregation.")
-            .allowableValues(ResultOutputStrategy.PER_RESPONSE.getValue(), ResultOutputStrategy.PER_HIT.getValue())
+            .allowableValues(ResultOutputStrategy.getNonPaginatedResponseOutputStrategies())
             .defaultValue(ResultOutputStrategy.PER_RESPONSE.getValue())
             .required(true)
             .expressionLanguageSupported(ExpressionLanguageScope.NONE)
diff --git a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/api/ResultOutputStrategy.java b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/api/ResultOutputStrategy.java
index 01c84cbdc1..5ee50f7369 100644
--- a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/api/ResultOutputStrategy.java
+++ b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/api/ResultOutputStrategy.java
@@ -20,6 +20,7 @@ package org.apache.nifi.processors.elasticsearch.api;
 import org.apache.nifi.components.DescribedValue;
 
 import java.util.Arrays;
+import java.util.EnumSet;
 
 public enum ResultOutputStrategy implements DescribedValue {
     PER_HIT("splitUp-yes", "Flowfile per hit."),
@@ -52,6 +53,10 @@ public enum ResultOutputStrategy implements DescribedValue {
         return description;
     }
 
+    public static EnumSet<ResultOutputStrategy> getNonPaginatedResponseOutputStrategies() {
+        return EnumSet.of(PER_RESPONSE, PER_HIT);
+    }
+
     public static ResultOutputStrategy fromValue(final String value) {
         return Arrays.stream(ResultOutputStrategy.values()).filter(v -> v.getValue().equals(value)).findFirst()
                 .orElseThrow(() -> new IllegalArgumentException(String.format("Unknown value %s", value)));
diff --git a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/test/groovy/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearchTest.groovy b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/test/groovy/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearchTest.groovy
index 25707cd300..670b9599b0 100644
--- a/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/test/groovy/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearchTest.groovy
+++ b/nifi-nar-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/test/groovy/org/apache/nifi/processors/elasticsearch/AbstractJsonQueryElasticsearchTest.groovy
@@ -21,14 +21,16 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import org.apache.nifi.components.state.Scope
 import org.apache.nifi.flowfile.FlowFile
 import org.apache.nifi.processors.elasticsearch.api.AggregationResultsFormat
-import org.apache.nifi.processors.elasticsearch.api.SearchResultsFormat
 import org.apache.nifi.processors.elasticsearch.api.ResultOutputStrategy
+import org.apache.nifi.processors.elasticsearch.api.SearchResultsFormat
 import org.apache.nifi.provenance.ProvenanceEventType
 import org.apache.nifi.util.MockFlowFile
 import org.apache.nifi.util.TestRunner
 import org.apache.nifi.util.TestRunners
 import org.junit.jupiter.api.Test
 
+import java.util.stream.Collectors
+
 import static groovy.json.JsonOutput.prettyPrint
 import static groovy.json.JsonOutput.toJson
 import static org.hamcrest.CoreMatchers.equalTo
@@ -95,9 +97,12 @@ abstract class AbstractJsonQueryElasticsearchTest<P extends AbstractJsonQueryEla
         runner.setProperty(AbstractJsonQueryElasticsearch.SEARCH_RESULTS_SPLIT, "not-enum2")
         runner.setProperty(AbstractJsonQueryElasticsearch.OUTPUT_NO_HITS, "not-boolean")
 
+        final String nonPaginatedResultOutputStrategies = ResultOutputStrategy.getNonPaginatedResponseOutputStrategies()
+                .stream().map(r -> r.getValue())
+                .collect(Collectors.joining(", "))
         final String expectedAllowedSplitHits = processor instanceof AbstractPaginatedJsonQueryElasticsearch
             ? ResultOutputStrategy.values().collect {r -> r.getValue()}.join(", ")
-            : [ResultOutputStrategy.PER_RESPONSE.getValue(), ResultOutputStrategy.PER_HIT.getValue()].join(", ")
+            : nonPaginatedResultOutputStrategies
 
         final AssertionError assertionError = assertThrows(AssertionError.class, runner.&run)
         assertThat(assertionError.getMessage(), equalTo(String.format("Processor has 8 validation failures:\n" +
@@ -115,7 +120,7 @@ abstract class AbstractJsonQueryElasticsearchTest<P extends AbstractJsonQueryEla
                 AbstractJsonQueryElasticsearch.TYPE.getName(), AbstractJsonQueryElasticsearch.TYPE.getName(),
                 AbstractJsonQueryElasticsearch.CLIENT_SERVICE.getDisplayName(),
                 AbstractJsonQueryElasticsearch.SEARCH_RESULTS_SPLIT.getName(), expectedAllowedSplitHits,
-                AbstractJsonQueryElasticsearch.AGGREGATION_RESULTS_SPLIT.getName(), [ResultOutputStrategy.PER_RESPONSE.getValue(), ResultOutputStrategy.PER_HIT.getValue()].join(", "),
+                AbstractJsonQueryElasticsearch.AGGREGATION_RESULTS_SPLIT.getName(), nonPaginatedResultOutputStrategies,
                 AbstractJsonQueryElasticsearch.OUTPUT_NO_HITS.getName(),
                 AbstractJsonQueryElasticsearch.CLIENT_SERVICE.getDisplayName()
         )))