am 00e20ef3: Support for primitive (non-Object) captors in littlemock

* commit '00e20ef3c702f7388c98bdeee300d218c7871f4d':
  Support for primitive (non-Object) captors in littlemock

2 files changed, 43 insertions, 0 deletions

diff --git a/src/com/google/testing/littlemock/ArgumentCaptor.java b/src/com/google/testing/littlemock/ArgumentCaptor.java
index 99e9b55..ff75246 100644
--- a/src/com/google/testing/littlemock/ArgumentCaptor.java
+++ b/src/com/google/testing/littlemock/ArgumentCaptor.java
@@ -34,4 +34,5 @@ public interface ArgumentCaptor<T extends Object> extends LittleMock.ArgumentMat
 
   /** Use this argument captor to perform the capture. */
   public T capture();
+
 }
diff --git a/src/com/google/testing/littlemock/LittleMock.java b/src/com/google/testing/littlemock/LittleMock.java
index ee589db..140f853 100644
--- a/src/com/google/testing/littlemock/LittleMock.java
+++ b/src/com/google/testing/littlemock/LittleMock.java
@@ -358,12 +358,27 @@ public class LittleMock {
     return new ArgumentCaptorImpl<T>();
   }
 
+  /**
+   * Create a correctly typed ArgumentCaptor that also works for primitive types.
+   * For example, to capture an int, use
+   * <code>ArgumentCaptor<Integer> c = createCaptor(Integer.class);</code>
+   */
+  public static <T> ArgumentCaptor<T> createCaptor(Class<T> clazz) {
+    return new ArgumentCaptorImpl<T>(clazz);
+  }
+
   /** Implementation of the {@link ArgumentCaptor} interface. */
   private static class ArgumentCaptorImpl<T extends Object> implements ArgumentCaptor<T> {
     private final ArrayList<T> mAllValues = new ArrayList<T>();
     private T mValue;
+    private Class<T> mClass;
 
     private ArgumentCaptorImpl() {
+        mClass = null;
+    }
+
+    private ArgumentCaptorImpl(Class<T> clazz) {
+        mClass = clazz;
     }
 
     public void setValue(T value) {
@@ -383,9 +398,36 @@ public class LittleMock {
 
     @Override
     public T capture() {
+      if (mClass != null) {
+        if (Integer.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Integer>addMatcher(this, 0);
+        }
+        if (Float.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Float>addMatcher(this, 0f);
+        }
+        if (Double.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Double>addMatcher(this, 0.0);
+        }
+        if (Boolean.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Boolean>addMatcher(this, false);
+        }
+        if (Short.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Short>addMatcher(this, (short) 0);
+        }
+        if (Character.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Character>addMatcher(this, '\u0000');
+        }
+        if (Long.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Long>addMatcher(this, 0L);
+        }
+        if (Byte.class.isAssignableFrom(mClass)) {
+          return (T) LittleMock.<Byte>addMatcher(this, (byte) 0);
+        }
+      }
       return LittleMock.<T>addMatcher(this, null);
     }
 
+
     @Override
     public boolean matches(Object value) {
       // A capture always matches any argument.