NIFI-4686 MockFlowFile now includes penalized status when copying. This closes #2438.

Signed-off-by: Mark Payne <markap14@hotmail.com>

nifi-mock/src/main/java/org/apache/nifi/util/MockFlowFile.java

@@ -94,6 +94,8 @@ public class MockFlowFile implements FlowFileRecord {
         final byte[] dataToCopy = ((MockFlowFile) toCopy).data;
         this.data = new byte[dataToCopy.length];
         System.arraycopy(dataToCopy, 0, this.data, 0, dataToCopy.length);
+
+        this.penalized = toCopy.isPenalized();
     }
 
     void setPenalized() {

nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java

@@ -89,6 +89,18 @@ public class TestMockProcessSession {
         session.transfer(ff1);
     }
 
+    @Test
+    public void testKeepPenalizedStatusAfterPuttingAttribute(){
+        final Processor processor = new PoorlyBehavedProcessor();
+        final MockProcessSession session = new MockProcessSession(new SharedSessionState(processor, new AtomicLong(0L)), processor);
+        FlowFile ff1 = session.createFlowFile("hello, world".getBytes());
+        ff1 = session.penalize(ff1);
+        assertEquals(true, ff1.isPenalized());
+        ff1 = session.putAttribute(ff1, "hello", "world");
+        // adding attribute to flow file should not override the original penalized status
+        assertEquals(true, ff1.isPenalized());
+    }
+
     protected static class PoorlyBehavedProcessor extends AbstractProcessor {
 
         private static final Relationship REL_FAILURE = new Relationship.Builder()

nifi-nar-bundles/nifi-couchbase-bundle/nifi-couchbase-processors/src/test/java/org/apache/nifi/processors/couchbase/TestGetCouchbaseKey.java

@@ -407,6 +407,7 @@ public class TestGetCouchbaseKey {
         MockFlowFile orgFile = testRunner.getFlowFilesForRelationship(REL_RETRY).get(0);
         orgFile.assertContentEquals(inputFileDataStr);
         orgFile.assertAttributeEquals(Exception.key(), exception.getClass().getName());
+        Assert.assertEquals(true, orgFile.isPenalized());
     }
 
     @Test