add simpler capture() method with lambda

mockito · Mar 8, 2016 · 9238c4c · 9238c4c
1 parent ce93783
commit 9238c4c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/mockito-kotlin/src/main/kotlin/com/nhaarman/mockito_kotlin/Mockito.kt b/mockito-kotlin/src/main/kotlin/com/nhaarman/mockito_kotlin/Mockito.kt
@@ -58,6 +58,10 @@ fun <T> Stubber.whenever(mock: T) = `when`(mock)
 
 inline fun <reified T : Any> argumentCaptor() = ArgumentCaptor.forClass(T::class.java)
 inline fun <reified T : Any> capture(captor: ArgumentCaptor<T>): T = captor.capture() ?: createInstance<T>()
+inline fun <reified T : Any> capture(noinline consumer: (T) -> Unit): T {
+    var times = 0
+    return argThat { if (++times == 1) consumer.invoke(this); true }
+}
 
 inline fun <reified T : Any> eq(value: T) = Mockito.eq(value) ?: createInstance<T>()
 inline fun <reified T : Any> anyArray(): Array<T> = Mockito.any(Array<T>::class.java) ?: arrayOf()

diff --git a/mockito-kotlin/src/test/kotlin/ArgumentCaptorTest.kt b/mockito-kotlin/src/test/kotlin/ArgumentCaptorTest.kt
@@ -8,7 +8,7 @@ import java.util.*
 
 class ArgumentCaptorTest {
     @Test
-    fun captor() {
+    fun explicitCaptor() {
         val date: Date = mock()
         val time = argumentCaptor<Long>()
 
@@ -17,4 +17,14 @@ class ArgumentCaptorTest {
         verify(date).time = capture(time)
         expect(time.value).toBe(5L)
     }
+
+    @Test
+    fun implicitCaptor() {
+        val date: Date = mock()
+        date.time = 5L
+
+        verify(date).time = capture {
+            expect(it).toBe(5L)
+        }
+    }
 }