From f58972e566448886cdf984e6c8a34c83674b0705 Mon Sep 17 00:00:00 2001
From: Mark Payne <markap14@hotmail.com>
Date: Thu, 25 Jun 2015 09:46:16 -0400
Subject: [PATCH] NIFI-728: Allow Mock Framework to use property descriptors
 from subclasses that are created for unit testing

---
 .../nifi/util/MockConfigurationContext.java    | 18 +++++++++++++++++-
 .../nifi/util/StandardProcessorTestRunner.java |  4 ++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockConfigurationContext.java b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockConfigurationContext.java
index 61af49da42..742f03be9c 100644
--- a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockConfigurationContext.java
+++ b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockConfigurationContext.java
@@ -22,14 +22,21 @@ import java.util.Map;
 import org.apache.nifi.components.PropertyDescriptor;
 import org.apache.nifi.components.PropertyValue;
 import org.apache.nifi.controller.ConfigurationContext;
+import org.apache.nifi.controller.ControllerService;
 import org.apache.nifi.controller.ControllerServiceLookup;
 
 public class MockConfigurationContext implements ConfigurationContext {
 
     private final Map<PropertyDescriptor, String> properties;
     private final ControllerServiceLookup serviceLookup;
+    private final ControllerService service;
 
     public MockConfigurationContext(final Map<PropertyDescriptor, String> properties, final ControllerServiceLookup serviceLookup) {
+        this(null, properties, serviceLookup);
+    }
+
+    public MockConfigurationContext(final ControllerService service, final Map<PropertyDescriptor, String> properties, final ControllerServiceLookup serviceLookup) {
+        this.service = service;
         this.properties = properties;
         this.serviceLookup = serviceLookup;
     }
@@ -38,7 +45,7 @@ public class MockConfigurationContext implements ConfigurationContext {
     public PropertyValue getProperty(final PropertyDescriptor property) {
         String value = properties.get(property);
         if (value == null) {
-            value = property.getDefaultValue();
+            value = getActualDescriptor(property).getDefaultValue();
         }
         return new MockPropertyValue(value, serviceLookup);
     }
@@ -47,4 +54,13 @@ public class MockConfigurationContext implements ConfigurationContext {
     public Map<PropertyDescriptor, String> getProperties() {
         return new HashMap<>(this.properties);
     }
+
+    private PropertyDescriptor getActualDescriptor(final PropertyDescriptor property) {
+        if (service == null) {
+            return property;
+        }
+
+        final PropertyDescriptor resolved = service.getPropertyDescriptor(property.getName());
+        return resolved == null ? property : resolved;
+    }
 }
diff --git a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/StandardProcessorTestRunner.java b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/StandardProcessorTestRunner.java
index 89385477d7..048e2b92c3 100644
--- a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/StandardProcessorTestRunner.java
+++ b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/StandardProcessorTestRunner.java
@@ -215,7 +215,7 @@ public class StandardProcessorTestRunner implements TestRunner {
             executorService.shutdown();
             try {
                 executorService.awaitTermination(runWait, TimeUnit.MILLISECONDS);
-            } catch (InterruptedException e1) {
+            } catch (final InterruptedException e1) {
             }
 
             int finishedCount = 0;
@@ -609,7 +609,7 @@ public class StandardProcessorTestRunner implements TestRunner {
         }
 
         try {
-            final ConfigurationContext configContext = new MockConfigurationContext(configuration.getProperties(), context);
+            final ConfigurationContext configContext = new MockConfigurationContext(service, configuration.getProperties(), context);
             ReflectionUtils.invokeMethodsWithAnnotation(OnEnabled.class, service, configContext);
         } catch (final InvocationTargetException ite) {
             ite.getCause().printStackTrace();