Support nested test classes

Author: Austin Brolly

src/main/java/com/github/exabrial/junit5/injectmap/InjectExtension.java

@@ -13,6 +13,7 @@
 
 import org.junit.jupiter.api.extension.BeforeTestExecutionCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.TestInstances;
 import org.mockito.InjectMocks;
 
 import javassist.util.proxy.MethodFilter;
@@ -37,31 +38,46 @@ public static boolean status() {
 
 	@Override
 	public void beforeTestExecution(ExtensionContext context) throws Exception {
-		Object testInstance = context.getTestInstance().get();
-		if (testInstance != null) {
-			final Map<String, Field> injectMap = new HashMap<>();
-			for (Field testClassField : testInstance.getClass().getDeclaredFields()) {
-				if (testClassField.getAnnotation(InjectMocks.class) != null) {
-					testClassField.setAccessible(true);
-					final Object injectionTarget = testClassField.get(testInstance);
-					final ProxyFactory proxyFactory = new ProxyFactory();
-					proxyFactory.setSuperclass(injectionTarget.getClass());
-					proxyFactory.setFilter(createMethodFilter());
-					final Class<?> proxyClass = proxyFactory.createClass();
-					final Object proxy = proxyClass.newInstance();
-					final Map<String, List<Field>> fieldMap = createFieldMap(injectionTarget.getClass());
-					Method postConstructMethod;
-					if (testClassField.getAnnotation(InvokePostConstruct.class) != null) {
-						postConstructMethod = findPostConstructMethod(injectionTarget);
-					} else {
-						postConstructMethod = null;
-					}
-					final MethodHandler handler = createMethodHandler(injectMap, injectionTarget, fieldMap, testInstance, postConstructMethod);
-					((Proxy) proxy).setHandler(handler);
-					testClassField.set(testInstance, proxy);
-				} else if (testClassField.getAnnotation(InjectionSource.class) != null) {
-					injectMap.put(testClassField.getName(), testClassField);
+		final Object testInstance = context.getTestInstance().get();
+		final Class<?> actualTestClazz = testInstance.getClass();
+		System.out.println("actualTestClass: " + actualTestClazz.getName());
+		Class<?> topLevelClass = actualTestClazz;
+		while (topLevelClass.getEnclosingClass() != null) {
+			topLevelClass = topLevelClass.getEnclosingClass();
+		}
+		final Class<?> parentTestClass = topLevelClass != actualTestClazz ? topLevelClass : null;
+
+		final Object parentClassInstance = context
+				.getTestInstances()
+				.get()
+				.getEnclosingInstances()
+				.stream()
+				.filter(i -> i.getClass() == parentTestClass)
+				.findFirst()
+				.orElse(testInstance);
+		System.out.println("enclosing instance class name: " + parentClassInstance.getClass().getName());
+		final Map<String, Field> injectMap = new HashMap<>();
+		for (Field testClassField : parentClassInstance.getClass().getDeclaredFields()) {
+			if (testClassField.getAnnotation(InjectMocks.class) != null) {
+				testClassField.setAccessible(true);
+				final Object injectionTarget = testClassField.get(parentClassInstance);
+				final ProxyFactory proxyFactory = new ProxyFactory();
+				proxyFactory.setSuperclass(injectionTarget.getClass());
+				proxyFactory.setFilter(createMethodFilter());
+				final Class<?> proxyClass = proxyFactory.createClass();
+				final Object proxy = proxyClass.newInstance();
+				final Map<String, List<Field>> fieldMap = createFieldMap(injectionTarget.getClass());
+				Method postConstructMethod;
+				if (testClassField.getAnnotation(InvokePostConstruct.class) != null) {
+					postConstructMethod = findPostConstructMethod(injectionTarget);
+				} else {
+					postConstructMethod = null;
 				}
+				final MethodHandler handler = createMethodHandler(injectMap, injectionTarget, fieldMap, parentClassInstance, postConstructMethod);
+				((Proxy) proxy).setHandler(handler);
+				testClassField.set(parentClassInstance, proxy);
+			} else if (testClassField.getAnnotation(InjectionSource.class) != null) {
+				injectMap.put(testClassField.getName(), testClassField);
 			}
 		}
 	}
@@ -73,21 +89,16 @@ private Method findPostConstructMethod(Object injectionTarget) {
 			}
 		}
 		throw new RuntimeException(
-				"@InvokePostConstruct is delcared on:" + injectionTarget + " however no method annotated with @PostConstruct found");
+				"@InvokePostConstruct is declared on:" + injectionTarget + " however no method annotated with @PostConstruct found");
 	}
 
-	private Map<String, List<Field>> createFieldMap(Class<? extends Object> targetClass) {
+	private Map<String, List<Field>> createFieldMap(Class<?> targetClass) {
 		if (targetClass == Object.class) {
 			return new HashMap<>();
 		} else {
 			Map<String, List<Field>> fieldMap = createFieldMap(targetClass.getSuperclass());
 			for (Field field : targetClass.getDeclaredFields()) {
-				List<Field> fieldList = fieldMap.get(field.getName());
-				if (fieldList == null) {
-					fieldList = new LinkedList<>();
-					fieldMap.put(field.getName(), fieldList);
-				}
-				fieldList.add(field);
+				fieldMap.computeIfAbsent(field.getName(), k -> new LinkedList<>()).add(field);
 			}
 			return fieldMap;
 		}