From 58f60b3cf0430e1ca5141c550595cd827b7e4b3e Mon Sep 17 00:00:00 2001 From: m-hogue Date: Thu, 6 Jul 2017 12:10:09 -0400 Subject: [PATCH] NIFI-3218: throw exception in MockProcessSession when transferring file to self --- .../java/org/apache/nifi/util/MockProcessSession.java | 7 +++++++ .../org/apache/nifi/util/TestMockProcessSession.java | 9 +++++++++ 2 files changed, 16 insertions(+) diff --git a/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java b/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java index af472d6448..553c6e4e0a 100644 --- a/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java +++ b/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java @@ -756,6 +756,13 @@ public class MockProcessSession implements ProcessSession { throw new IllegalArgumentException("I only accept MockFlowFile"); } + // if the flowfile provided was created in this session (i.e. it's in currentVersions), + // then throw an exception indicating that you can't transfer flowfiles back to self. + // this mimics the behavior of StandardProcessSession + if(currentVersions.get(flowFile.getId()) != null) { + throw new IllegalArgumentException("Cannot transfer FlowFiles that are created in this Session back to self"); + } + final MockFlowFile mockFlowFile = (MockFlowFile) flowFile; beingProcessed.remove(flowFile.getId()); processorQueue.offer(mockFlowFile); diff --git a/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java b/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java index 1c554d2007..7173f72aa2 100644 --- a/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java +++ b/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java @@ -80,6 +80,15 @@ public class TestMockProcessSession { } + @Test(expected = IllegalArgumentException.class) + public void testRejectTransferNewlyCreatedFileToSelf() { + final Processor processor = new PoorlyBehavedProcessor(); + final MockProcessSession session = new MockProcessSession(new SharedSessionState(processor, new AtomicLong(0L)), processor); + final FlowFile ff1 = session.createFlowFile("hello, world".getBytes()); + // this should throw an exception because we shouldn't allow a newly created flowfile to get routed back to self + session.transfer(ff1); + } + protected static class PoorlyBehavedProcessor extends AbstractProcessor { private static final Relationship REL_FAILURE = new Relationship.Builder()