From 96a42cd72b0f009377080d471c420e14d23eca89 Mon Sep 17 00:00:00 2001
From: Mark Payne <markap14@hotmail.com>
Date: Sun, 1 Mar 2015 15:30:42 -0500
Subject: [PATCH] NIFI-389: Fixed bug that caused services that depend on other
 services to not be locatable in mock framework

---
 .../org/apache/nifi/util/MockControllerServiceLookup.java     | 3 +++
 .../main/java/org/apache/nifi/util/MockProcessContext.java    | 4 +++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockControllerServiceLookup.java b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockControllerServiceLookup.java
index f5a716c99f..8298a39617 100644
--- a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockControllerServiceLookup.java
+++ b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockControllerServiceLookup.java
@@ -47,6 +47,9 @@ public abstract class MockControllerServiceLookup implements ControllerServiceLo
         controllerServiceMap.remove(service.getIdentifier());
     }
 
+    protected void addControllerServices(final MockControllerServiceLookup other) {
+        this.controllerServiceMap.putAll(other.controllerServiceMap);
+    }
     
     protected ControllerServiceConfiguration getConfiguration(final String identifier) {
         return controllerServiceMap.get(identifier);
diff --git a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessContext.java b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessContext.java
index f83301c107..65369288e7 100644
--- a/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessContext.java
+++ b/nifi/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessContext.java
@@ -67,11 +67,13 @@ public class MockProcessContext extends MockControllerServiceLookup implements S
             annotationData = context.getControllerServiceAnnotationData(component);
             final Map<PropertyDescriptor, String> props = context.getControllerServiceProperties(component);
             properties.putAll(props);
+
+            super.addControllerServices(context);
         } catch (IllegalArgumentException e) {
             // do nothing...the service is being loaded
         }
     }
-
+    
     @Override
     public PropertyValue getProperty(final PropertyDescriptor descriptor) {
         return getProperty(descriptor.getName());