From 0c4b9e56e1729723039339ceee4dc78e0cd74862 Mon Sep 17 00:00:00 2001
From: Matt Burgess <mattyb149@apache.org>
Date: Thu, 13 Oct 2016 09:50:56 -0400
Subject: [PATCH] NIFI-2819: Added support for Expresssion Language in
 ModifyBytes

---
 .../nifi/processors/standard/ModifyBytes.java |  6 ++--
 .../processors/standard/TestModifyBytes.java  | 35 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ModifyBytes.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ModifyBytes.java
index 444a1729dee8..6d8bee323654 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ModifyBytes.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ModifyBytes.java
@@ -66,6 +66,7 @@ public class ModifyBytes extends AbstractProcessor {
             .required(true)
             .addValidator(StandardValidators.DATA_SIZE_VALIDATOR)
             .defaultValue("0 B")
+            .expressionLanguageSupported(true)
             .build();
     public static final PropertyDescriptor END_OFFSET = new PropertyDescriptor.Builder()
             .name("End Offset")
@@ -73,6 +74,7 @@ public class ModifyBytes extends AbstractProcessor {
             .required(true)
             .addValidator(StandardValidators.DATA_SIZE_VALIDATOR)
             .defaultValue("0 B")
+            .expressionLanguageSupported(true)
             .build();
     public static final PropertyDescriptor REMOVE_ALL = new PropertyDescriptor.Builder()
             .name("Remove All Content")
@@ -114,8 +116,8 @@ public void onTrigger(ProcessContext context, ProcessSession session) throws Pro
 
         final ComponentLog logger = getLogger();
 
-        final long startOffset = context.getProperty(START_OFFSET).asDataSize(DataUnit.B).longValue();
-        final long endOffset = context.getProperty(END_OFFSET).asDataSize(DataUnit.B).longValue();
+        final long startOffset = context.getProperty(START_OFFSET).evaluateAttributeExpressions(ff).asDataSize(DataUnit.B).longValue();
+        final long endOffset = context.getProperty(END_OFFSET).evaluateAttributeExpressions(ff).asDataSize(DataUnit.B).longValue();
         final boolean removeAll = context.getProperty(REMOVE_ALL).asBoolean();
         final long newFileSize = removeAll ? 0L : ff.getSize() - startOffset - endOffset;
 
diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestModifyBytes.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestModifyBytes.java
index 7243be867634..91f08531b575 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestModifyBytes.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestModifyBytes.java
@@ -21,6 +21,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.HashMap;
 
 import org.apache.nifi.util.MockFlowFile;
 import org.apache.nifi.util.TestRunner;
@@ -94,6 +95,24 @@ public void testRemoveHeader() throws IOException {
         out.assertContentEquals(noHeaderFile);
     }
 
+    @Test
+    public void testRemoveHeaderEL() throws IOException {
+        final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());
+        runner.setProperty(ModifyBytes.START_OFFSET, "${numBytes}"); //REMOVE - '<<<HEADER>>>'
+        runner.setProperty(ModifyBytes.END_OFFSET, "0 MB");
+
+        runner.enqueue(testFilePath, new HashMap<String, String>() {{
+            put("numBytes", "12 B");
+        }});
+        runner.run();
+
+        runner.assertAllFlowFilesTransferred(ModifyBytes.REL_SUCCESS, 1);
+        final MockFlowFile out = runner.getFlowFilesForRelationship(ModifyBytes.REL_SUCCESS).get(0);
+        final String outContent = new String(out.toByteArray(), StandardCharsets.UTF_8);
+        System.out.println(outContent);
+        out.assertContentEquals(noHeaderFile);
+    }
+
     @Test
     public void testKeepFooter() throws IOException {
         final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());
@@ -124,6 +143,22 @@ public void testKeepHeader() throws IOException {
         out.assertContentEquals("<<<HEADER>>>".getBytes("UTF-8"));
     }
 
+    @Test
+    public void testKeepHeaderEL() throws IOException {
+        final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());
+        runner.setProperty(ModifyBytes.START_OFFSET, "0 B");
+        runner.setProperty(ModifyBytes.END_OFFSET, "${numBytes}");
+
+        runner.enqueue(testFilePath, new HashMap<String, String>() {{
+            put("numBytes", "181 B");
+        }});
+        runner.run();
+
+        runner.assertAllFlowFilesTransferred(ModifyBytes.REL_SUCCESS, 1);
+        final MockFlowFile out = runner.getFlowFilesForRelationship(ModifyBytes.REL_SUCCESS).get(0);
+        out.assertContentEquals("<<<HEADER>>>".getBytes("UTF-8"));
+    }
+
     @Test
     public void testRemoveFooter() throws IOException {
         final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());