1 files changed, 22 insertions, 12 deletions
diff --git a/src/com/google/testing/littlemock/LittleMock.java b/src/com/google/testing/littlemock/LittleMock.java
index c8f2b2f..efd205c 100644
--- a/src/com/google/testing/littlemock/LittleMock.java
+++ b/src/com/google/testing/littlemock/LittleMock.java
@@ -1011,18 +1011,6 @@ public class LittleMock {
     return DEFAULT_RETURN_VALUE_LOOKUP.get(returnType);
   }
 
-  /**
-   * If the input object is one of our mocks, returns the {@link DefaultInvocationHandler}
-   * we constructed it with.  Otherwise fails with {@link IllegalArgumentException}.
-   */
-  private static DefaultInvocationHandler getHandlerFrom(Object mock) {
-    InvocationHandler invocationHandler = Proxy.getInvocationHandler(mock);
-    if (!(invocationHandler instanceof DefaultInvocationHandler)) {
-      throw new IllegalArgumentException("not a valid mock: " + mock);
-    }
-    return (DefaultInvocationHandler) invocationHandler;
-  }
-
   /** Gets a suitable class loader for use with the proxy. */
   private static ClassLoader getClassLoader() {
     return LittleMock.class.getClassLoader();
@@ -1042,6 +1030,28 @@ public class LittleMock {
     }
   }
 
+  /**
+   * If the input object is one of our mocks, returns the {@link DefaultInvocationHandler}
+   * we constructed it with.  Otherwise fails with {@link IllegalArgumentException}.
+   */
+  private static DefaultInvocationHandler getHandlerFrom(Object mock) {
+    try {
+      InvocationHandler invocationHandler = Proxy.getInvocationHandler(mock);
+      if (invocationHandler instanceof DefaultInvocationHandler) {
+        return (DefaultInvocationHandler) invocationHandler;
+      }
+    } catch (IllegalArgumentException expectedIfNotAProxy) {}
+    try {
+      Class<?> proxyBuilder = Class.forName("com.google.dexmaker.stock.ProxyBuilder");
+      Method getHandlerMethod = proxyBuilder.getMethod("getInvocationHandler", Object.class);
+      Object invocationHandler = getHandlerMethod.invoke(proxyBuilder, mock);
+      if (invocationHandler instanceof DefaultInvocationHandler) {
+        return (DefaultInvocationHandler) invocationHandler;
+      }
+    } catch (Exception expectedIfNotAProxyBuilderMock) {}
+    throw new IllegalArgumentException("not a valid mock: " + mock);
+  }
+
   /** Create a dynamic proxy for the given class, delegating to the given invocation handler. */
   private static Object createProxy(Class<?> clazz, InvocationHandler handler) {
     if (clazz.isInterface()) {