[0.1.38] Improve support for mocking Mono/Flux return value methods

unloggedio · Dec 28, 2023 · ba47198 · ba47198
1 parent 2fcd670
commit ba47198
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 33 deletions.
diff --git a/pom.xml b/pom.xml
@@ -10,7 +10,7 @@
 
     <artifactId>unlogged-sdk</artifactId>
     <groupId>video.bug</groupId>
-    <version>0.1.37</version>
+    <version>0.1.38</version>
 
     <properties>
         <maven.compiler.source>8</maven.compiler.source>

diff --git a/src/main/java/io/unlogged/AgentCommandExecutorImpl.java b/src/main/java/io/unlogged/AgentCommandExecutorImpl.java
@@ -393,7 +393,7 @@ public AgentCommandRawResponse executeCommandRaw(AgentCommandRequest agentComman
             } catch (Throwable exception) {
                 if (exception instanceof InvocationTargetException) {
                     agentCommandResponse.setResponseType(ResponseType.EXCEPTION);
-//                    exception.getCause().printStackTrace();
+                    exception.getCause().printStackTrace();
                 } else {
                     agentCommandResponse.setResponseType(ResponseType.FAILED);
 //                    exception.printStackTrace();
@@ -662,34 +662,40 @@ public Object arrangeMocks(Class<?> targetClassType, ClassLoader targetClassLoad
                 .collect(Collectors.groupingBy(DeclaredMock::getFieldName,
                         Collectors.toList()));
 
-        DynamicType.Unloaded<? extends Object> extendedClass = byteBuddyInstance
-                .subclass(targetClassType)
-                .make();
 
-        ClassLoadingStrategy<ClassLoader> strategy;
-        if (ClassInjector.UsingLookup.isAvailable()) {
-            Class<?> methodHandles = null;
-            try {
-                methodHandles = targetClassLoader.loadClass("java.lang.invoke.MethodHandles");
-                Object lookup = methodHandles.getMethod("lookup").invoke(null);
-                Method privateLookupIn = methodHandles.getMethod("privateLookupIn",
-                        Class.class,
-                        targetClassLoader.loadClass("java.lang.invoke.MethodHandles$Lookup"));
-                Object privateLookup = privateLookupIn.invoke(null, targetClassType, lookup);
-                strategy = ClassLoadingStrategy.UsingLookup.of(privateLookup);
-            } catch (Exception e) {
-                // should never happen
-                throw new RuntimeException(e);
-            }
-        } else if (ClassInjector.UsingReflection.isAvailable()) {
-            strategy = ClassLoadingStrategy.Default.INJECTION;
+        Object extendedClassInstance;
+        if (Modifier.isFinal(targetClassType.getModifiers())) {
+            extendedClassInstance = objectInstanceByClass;
         } else {
-            throw new IllegalStateException("No code generation strategy available");
-        }
+            DynamicType.Unloaded<? extends Object> extendedClass = byteBuddyInstance
+                    .subclass(targetClassType)
+                    .make();
 
-        DynamicType.Loaded<? extends Object> extendedClassLoaded = extendedClass.load(targetClassLoader, strategy);
+            ClassLoadingStrategy<ClassLoader> strategy;
+            if (ClassInjector.UsingLookup.isAvailable()) {
+                Class<?> methodHandles = null;
+                try {
+                    methodHandles = targetClassLoader.loadClass("java.lang.invoke.MethodHandles");
+                    Object lookup = methodHandles.getMethod("lookup").invoke(null);
+                    Method privateLookupIn = methodHandles.getMethod("privateLookupIn",
+                            Class.class,
+                            targetClassLoader.loadClass("java.lang.invoke.MethodHandles$Lookup"));
+                    Object privateLookup = privateLookupIn.invoke(null, targetClassType, lookup);
+                    strategy = ClassLoadingStrategy.UsingLookup.of(privateLookup);
+                } catch (Exception e) {
+                    // should never happen
+                    throw new RuntimeException(e);
+                }
+            } else if (ClassInjector.UsingReflection.isAvailable()) {
+                strategy = ClassLoadingStrategy.Default.INJECTION;
+            } else {
+                throw new IllegalStateException("No code generation strategy available");
+            }
 
-        Object extendedClassInstance = objenesis.newInstance(extendedClassLoaded.getLoaded());
+            DynamicType.Loaded<? extends Object> extendedClassLoaded = extendedClass.load(targetClassLoader, strategy);
+
+            extendedClassInstance = objenesis.newInstance(extendedClassLoaded.getLoaded());
+        }
 
 
         Class<?> fieldCopyForClass = targetClassType;
@@ -739,9 +745,10 @@ public Object arrangeMocks(Class<?> targetClassType, ClassLoader targetClassLoad
                                         field.getType()
                                 );
                             }
-
+                            String fieldTypeName = declaredMocksForField.get(0).getFieldTypeName();
+                            Class<?> classTypeToBeMocked = Class.forName(fieldTypeName);
                             existingMockInstance = createMockedInstance(targetClassLoader, objectInstanceByClass,
-                                    field, declaredMocksForField, fieldValue, field.getType());
+                                    field, declaredMocksForField, fieldValue, classTypeToBeMocked);
                             mockInstanceMap.put(key, existingMockInstance);
 
                         } else {

diff --git a/src/main/java/io/unlogged/Constants.java b/src/main/java/io/unlogged/Constants.java
@@ -5,5 +5,5 @@ public class Constants {
 
     private Constants() {}
 
-    public static final String AGENT_VERSION = "0.1.37";
+    public static final String AGENT_VERSION = "0.1.38";
 }
diff --git a/src/main/java/io/unlogged/mocking/MockHandler.java b/src/main/java/io/unlogged/mocking/MockHandler.java
@@ -212,11 +212,21 @@ private Object createReturnValueInstance(ThenParameter thenParameter, Method inv
                 }
 
                 if (typeReference.getRawClass().getCanonicalName().equals("reactor.core.publisher.Mono")) {
-                    returnValueInstance = Mono.just(jsonDeserializer.createInstance(returnValueSerialized,
-                            typeReference.containedType(0)));
+                    if (returnValueSerialized.startsWith("[")) {
+                        returnValueInstance = Mono.just(jsonDeserializer.createInstance(returnValueSerialized,
+                                typeFactory.constructArrayType(typeReference.containedType(0))));
+                    } else  {
+                        returnValueInstance = Mono.just(jsonDeserializer.createInstance(returnValueSerialized,
+                                typeReference.containedType(0)));
+                    }
                 } else if (typeReference.getRawClass().getCanonicalName().equals("reactor.core.publisher.Flux")) {
-                    returnValueInstance = Flux.just(jsonDeserializer.createInstance(returnValueSerialized,
-                            typeFactory.constructArrayType(typeReference.containedType(0))));
+                    if (returnValueSerialized.startsWith("[")) {
+                        returnValueInstance = Flux.just(jsonDeserializer.createInstance(returnValueSerialized,
+                                typeFactory.constructArrayType(typeReference.containedType(0))));
+                    } else {
+                        returnValueInstance = Flux.just(jsonDeserializer.createInstance(returnValueSerialized,
+                                typeReference.containedType(0)));
+                    }
                 } else {
                     returnValueInstance = jsonDeserializer.createInstance(returnValueSerialized, typeReference);
                 }