#20 Refactor mock/spy handling with DBuilderExtn

Provides a cleaner separation for handling mocks (supplied beans) and spy (enhancement)

Author: rbygrave

src/main/java/io/dinject/BootContext.java

@@ -3,7 +3,7 @@
 import io.dinject.core.BeanContextFactory;
 import io.dinject.core.Builder;
 import io.dinject.core.BuilderFactory;
-import io.dinject.core.SpyConsumer;
+import io.dinject.core.EnrichBean;
 import io.dinject.core.SuppliedBean;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -56,7 +56,7 @@ public class BootContext {
 
   private final List<SuppliedBean> suppliedBeans = new ArrayList<>();
 
-  private final List<SpyConsumer> spyConsumers = new ArrayList<>();
+  private final List<EnrichBean> enrichBeans = new ArrayList<>();
 
   private final Set<String> includeModules = new LinkedHashSet<>();
 
@@ -227,15 +227,15 @@ public <D> BootContext withMock(Class<D> type, Consumer<D> consumer) {
    * Use a mockito spy when injecting this bean type.
    */
   public <D> BootContext withSpy(Class<D> type) {
-    spyConsumers.add(new SpyConsumer<>(type, null));
+    enrichBeans.add(new EnrichBean<>(type, null));
     return this;
   }
 
   /**
    * Use a mockito spy when injecting this bean type additionally running setup on the spy instance.
    */
   public <D> BootContext withSpy(Class<D> type, Consumer<D> consumer) {
-    spyConsumers.add(new SpyConsumer<>(type, consumer));
+    enrichBeans.add(new EnrichBean<>(type, consumer));
     return this;
   }
 
@@ -251,7 +251,7 @@ public BeanContext load() {
     Set<String> moduleNames = factoryOrder.orderFactories();
     log.debug("building context with modules {}", moduleNames);
 
-    Builder rootBuilder = BuilderFactory.newRootBuilder(suppliedBeans, spyConsumers);
+    Builder rootBuilder = BuilderFactory.newRootBuilder(suppliedBeans, enrichBeans);
 
     for (BeanContextFactory factory : factoryOrder.factories()) {
       rootBuilder.addChild(factory.createContext(rootBuilder));

src/main/java/io/dinject/core/Builder.java

@@ -126,11 +126,12 @@ public interface Builder {
   BeanContext build();
 
   /**
-   * Return a spy enhanced bean for registration into the context.
+   * Return a potentially enriched bean for registration into the context.
+   * Typically for use with mockito spy.
    *
    * @param bean  The bean with dependencies injected
    * @param types The types this bean registers for
-   * @return Either the bean or the spy enhanced bean to register into the context.
+   * @return Either the bean or the enriched bean to register into the context.
    */
-  Object spy(Object bean, Class<?>[] types);
+  Object enrich(Object bean, Class<?>[] types);
 }

src/main/java/io/dinject/core/BuilderFactory.java

@@ -13,10 +13,15 @@ public class BuilderFactory {
    * Create the root level Builder.
    *
    * @param suppliedBeans The list of beans (typically test doubles) supplied when building the context.
-   * @param spyConsumers The list of classes we want to have with mockito spy enhancement
+   * @param enrichBeans   The list of classes we want to have with mockito spy enhancement
    */
-  public static Builder newRootBuilder(List<SuppliedBean> suppliedBeans, List<SpyConsumer> spyConsumers) {
-    return new DBuilder(suppliedBeans, spyConsumers);
+  public static Builder newRootBuilder(List<SuppliedBean> suppliedBeans, List<EnrichBean> enrichBeans) {
+
+    if (suppliedBeans.isEmpty() && enrichBeans.isEmpty()) {
+      // simple case, no mocks or spies
+      return new DBuilder();
+    }
+    return new DBuilderExtn(suppliedBeans, enrichBeans);
   }
 
   /**

src/main/java/io/dinject/core/DBuilder.java

@@ -6,7 +6,6 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,12 +29,7 @@ class DBuilder implements Builder {
   /**
    * The beans created and added to the context during building.
    */
-  private final DBeanMap beanMap = new DBeanMap();
-
-  /**
-   * Supplied beans (test doubles) given to the context prior to building.
-   */
-  private final boolean hasSuppliedBeans;
+  final DBeanMap beanMap = new DBeanMap();
 
   /**
    * The context/module name.
@@ -54,14 +48,12 @@ class DBuilder implements Builder {
 
   private final Map<String, BeanContext> children = new LinkedHashMap<>();
 
-  private final Map<Class<?>, SpyConsumer> spyMap;
-
   /**
    * Debug of the current bean being wired - used in injection errors.
    */
   private Class<?> injectTarget;
 
-  private Builder parent;
+  Builder parent;
 
   /**
    * Create a named context for non-root builders.
@@ -70,29 +62,15 @@ class DBuilder implements Builder {
     this.name = name;
     this.provides = provides;
     this.dependsOn = dependsOn;
-    this.hasSuppliedBeans = false;
-    this.spyMap = null;
   }
 
   /**
-   * Create for the root builder with supplied beans (test doubles).
+   * Create for the root builder.
    */
-  DBuilder(List<SuppliedBean> suppliedBeans, List<SpyConsumer> spyConsumers) {
+  DBuilder() {
     this.name = null;
     this.provides = null;
     this.dependsOn = null;
-    this.hasSuppliedBeans = (suppliedBeans != null && !suppliedBeans.isEmpty());
-    if (hasSuppliedBeans) {
-      beanMap.add(suppliedBeans);
-    }
-    if (spyConsumers == null || spyConsumers.isEmpty()) {
-      spyMap = null;
-    } else {
-      spyMap = new HashMap<>();
-      for (SpyConsumer spy : spyConsumers) {
-        spyMap.put(spy.getType(), spy);
-      }
-    }
   }
 
   @Override
@@ -117,9 +95,6 @@ public void setParent(Builder parent) {
 
   @Override
   public boolean isAddBeanFor(Class<?> addForType, Class<?> injectTarget) {
-    if (hasSuppliedBeans) {
-      return !beanMap.isSupplied(addForType.getName());
-    }
     if (parent == null) {
       return true;
     }
@@ -177,10 +152,8 @@ public void addChild(BeanContext child) {
   @Override
   public void register(Object bean, String name, Class<?>... types) {
     if (parent != null) {
-      // spy consumers only exist on top level builder
-      bean = parent.spy(bean, types);
-    } else {
-      bean = spy(bean, types);
+      // enrichment only exist on top level builder
+      bean = parent.enrich(bean, types);
     }
     beanMap.register(bean, name, types);
   }
@@ -189,27 +162,11 @@ public void register(Object bean, String name, Class<?>... types) {
    * Return the bean to register potentially with spy enhancement.
    */
   @Override
-  public Object spy(Object bean, Class<?>[] types) {
-    if (spyMap != null) {
-      SpyConsumer spyConsumer = spyMap.get(typeOf(bean, types));
-      if (spyConsumer != null) {
-        // enrich/enhance the bean for spying
-        return spyConsumer.spy(bean);
-      }
-    }
+  public Object enrich(Object bean, Class<?>[] types) {
+    // only enriched by DBuilderExtn
     return bean;
   }
 
-  /**
-   * Return the type to lookup for spy.
-   */
-  private Class<?> typeOf(Object bean, Class<?>... types) {
-    if (types.length > 0) {
-      return types[0];
-    }
-    return bean.getClass();
-  }
-
   @Override
   public void registerPrimary(Object bean, String name, Class<?>... types) {
     beanMap.registerPrimary(bean, name, types);

src/main/java/io/dinject/core/DBuilderExtn.java

@@ -0,0 +1,72 @@
+package io.dinject.core;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Extended builder that supports supplied beans (mocks) and enriching beans (spy).
+ */
+public class DBuilderExtn extends DBuilder {
+
+  private final Map<Class<?>, EnrichBean> enrichMap = new HashMap<>();
+
+  private final boolean hasSuppliedBeans;
+
+  DBuilderExtn(List<SuppliedBean> suppliedBeans, List<EnrichBean> enrichBeans) {
+    super();
+    this.hasSuppliedBeans = (suppliedBeans != null && !suppliedBeans.isEmpty());
+    if (hasSuppliedBeans) {
+      beanMap.add(suppliedBeans);
+    }
+    if (enrichBeans != null && !enrichBeans.isEmpty()) {
+      for (EnrichBean spy : enrichBeans) {
+        enrichMap.put(spy.getType(), spy);
+      }
+    }
+  }
+
+  /**
+   * Checking for supplied beans (mocks).
+   */
+  @Override
+  public boolean isAddBeanFor(Class<?> addForType, Class<?> injectTarget) {
+    if (hasSuppliedBeans) {
+      return !beanMap.isSupplied(addForType.getName());
+    }
+    return true;
+  }
+
+  /**
+   * Register beans with potential enhancement (spy).
+   */
+  @Override
+  public void register(Object bean, String name, Class<?>... types) {
+    if (parent != null) {
+      throw new IllegalStateException("no");
+    }
+    // this is the top level builder (parent always null)
+    bean = enrich(bean, types);
+    beanMap.register(bean, name, types);
+  }
+
+  /**
+   * Potentially enrich the bean prior to registering with context.
+   */
+  @SuppressWarnings("unchecked")
+  @Override
+  public Object enrich(Object bean, Class<?>[] types) {
+    EnrichBean enrich = enrichMap.get(typeOf(bean, types));
+    return enrich != null ? enrich.enrich(bean) : bean;
+  }
+
+  /**
+   * Return the type to lookup for enrichment.
+   */
+  private Class<?> typeOf(Object bean, Class<?>... types) {
+    if (types.length > 0) {
+      return types[0];
+    }
+    return bean.getClass();
+  }
+}

src/main/java/io/dinject/core/SpyConsumer.java renamed to src/main/java/io/dinject/core/EnrichBean.java

@@ -7,13 +7,13 @@
 /**
  * Holds Spy setup consumers for dependency injection using Mockito Spy.
  */
-public class SpyConsumer<B> {
+public class EnrichBean<B> {
 
   private final Class<B> type;
 
   private final Consumer<B> consumer;
 
-  public SpyConsumer(Class<B> type, Consumer<B> consumer) {
+  public EnrichBean(Class<B> type, Consumer<B> consumer) {
     this.type = type;
     this.consumer = consumer;
   }
@@ -28,8 +28,9 @@ public Class<B> getType() {
   /**
    * Return the spy enhanced bean instance to use.
    */
-  public B spy(B bean) {
+  public B enrich(B bean) {
 
+    // should extract a SPI for this. Only enrichment is Mockito spy at this point.
     B spy = Mockito.spy(bean);
     if (consumer != null) {
       consumer.accept(spy);

src/main/java/io/dinject/core/SuppliedBean.java

@@ -48,6 +48,7 @@ public Class<B> getType() {
    */
   public B getBean() {
     if (bean == null) {
+      // should extract a SPI for this
       bean = Mockito.mock(type);
     }
     if (consumer != null) {