From 8305da83c21061dbaf32cabefa83c16ba1a77e08 Mon Sep 17 00:00:00 2001
From: Stuart Douglas <stuart.w.douglas@gmail.com>
Date: Tue, 16 Feb 2021 09:32:28 +1100
Subject: [PATCH] Make sure mock is of correct type

Fixes #14991
---
 .../arc/processor/ClientProxyGenerator.java       | 15 +++++++++------
 .../java/io/quarkus/test/junit/MockSupport.java   |  4 ++++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/independent-projects/arc/processor/src/main/java/io/quarkus/arc/processor/ClientProxyGenerator.java b/independent-projects/arc/processor/src/main/java/io/quarkus/arc/processor/ClientProxyGenerator.java
index 7e8718e6295b4..4410f808f41d8 100644
--- a/independent-projects/arc/processor/src/main/java/io/quarkus/arc/processor/ClientProxyGenerator.java
+++ b/independent-projects/arc/processor/src/main/java/io/quarkus/arc/processor/ClientProxyGenerator.java
@@ -118,7 +118,7 @@ Collection<Resource> generate(BeanInfo bean, String beanClassName,
         FieldCreator beanField = clientProxy.getFieldCreator(BEAN_FIELD, DescriptorUtils.extToInt(beanClassName))
                 .setModifiers(ACC_PRIVATE | ACC_FINAL);
         if (mockable) {
-            clientProxy.getFieldCreator(MOCK_FIELD, Object.class).setModifiers(ACC_PRIVATE | ACC_VOLATILE);
+            clientProxy.getFieldCreator(MOCK_FIELD, providerType.descriptorName()).setModifiers(ACC_PRIVATE | ACC_VOLATILE);
         }
         FieldCreator contextField = null;
         if (BuiltinScope.APPLICATION.is(bean.getScope())) {
@@ -133,7 +133,7 @@ Collection<Resource> generate(BeanInfo bean, String beanClassName,
         implementGetContextualInstance(clientProxy, providerType);
         implementGetBean(clientProxy, beanField.getFieldDescriptor());
         if (mockable) {
-            implementMockMethods(clientProxy);
+            implementMockMethods(clientProxy, providerType);
         }
 
         for (MethodInfo method : getDelegatingMethods(bean, bytecodeTransformerConsumer, transformUnproxyableClasses)) {
@@ -231,17 +231,19 @@ Collection<Resource> generate(BeanInfo bean, String beanClassName,
         return classOutput.getResources();
     }
 
-    private void implementMockMethods(ClassCreator clientProxy) {
+    private void implementMockMethods(ClassCreator clientProxy, ProviderType providerType) {
         MethodCreator clear = clientProxy
                 .getMethodCreator(MethodDescriptor.ofMethod(clientProxy.getClassName(), CLEAR_MOCK_METHOD_NAME, void.class));
-        clear.writeInstanceField(FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, Object.class), clear.getThis(),
+        clear.writeInstanceField(FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, providerType.descriptorName()),
+                clear.getThis(),
                 clear.loadNull());
         clear.returnValue(null);
 
         MethodCreator set = clientProxy
                 .getMethodCreator(
                         MethodDescriptor.ofMethod(clientProxy.getClassName(), SET_MOCK_METHOD_NAME, void.class, Object.class));
-        set.writeInstanceField(FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, Object.class), set.getThis(),
+        set.writeInstanceField(FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, providerType.descriptorName()),
+                set.getThis(),
                 set.getMethodParam(0));
         set.returnValue(null);
     }
@@ -268,7 +270,8 @@ void implementDelegate(ClassCreator clientProxy, ProviderType providerType, Fiel
         if (mockable) {
             //if mockable and mocked just return the mock
             ResultHandle mock = creator.readInstanceField(
-                    FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, Object.class.getName()), creator.getThis());
+                    FieldDescriptor.of(clientProxy.getClassName(), MOCK_FIELD, providerType.descriptorName()),
+                    creator.getThis());
             BytecodeCreator falseBranch = creator.ifNull(mock).falseBranch();
             falseBranch.returnValue(falseBranch.checkCast(mock, providerType.className()));
         }
diff --git a/test-framework/junit5/src/main/java/io/quarkus/test/junit/MockSupport.java b/test-framework/junit5/src/main/java/io/quarkus/test/junit/MockSupport.java
index 2bbf9b6c1de55..60107393d702d 100644
--- a/test-framework/junit5/src/main/java/io/quarkus/test/junit/MockSupport.java
+++ b/test-framework/junit5/src/main/java/io/quarkus/test/junit/MockSupport.java
@@ -36,6 +36,10 @@ static <T> void installMock(T instance, T mock) {
         if (inst == null) {
             throw new IllegalStateException("No test in progress");
         }
+        if (!(instance.getClass().getSuperclass().isAssignableFrom(mock.getClass()))) {
+            throw new RuntimeException(mock
+                    + " is not assignable to type " + instance.getClass().getSuperclass());
+        }
         try {
             Method setMethod = instance.getClass().getDeclaredMethod("arc$setMock", Object.class);
             setMethod.invoke(instance, mock);