From a12e11f9c880d83e1e0626de01f842a7ccc8be1c Mon Sep 17 00:00:00 2001
From: igor <panel@go2.pl>
Date: Sun, 18 Feb 2018 18:48:17 +0000
Subject: [PATCH] Added argument captor overloads

---
 .../nhaarman/mockitokotlin2/ArgumentCaptor.kt | 14 ++++++++
 .../test/kotlin/test/ArgumentCaptorTest.kt    | 33 +++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/mockito-kotlin/src/main/kotlin/com/nhaarman/mockitokotlin2/ArgumentCaptor.kt b/mockito-kotlin/src/main/kotlin/com/nhaarman/mockitokotlin2/ArgumentCaptor.kt
index bc26bf51..9542446f 100644
--- a/mockito-kotlin/src/main/kotlin/com/nhaarman/mockitokotlin2/ArgumentCaptor.kt
+++ b/mockito-kotlin/src/main/kotlin/com/nhaarman/mockitokotlin2/ArgumentCaptor.kt
@@ -36,6 +36,13 @@ inline fun <reified T : Any> argumentCaptor(): KArgumentCaptor<T> {
     return KArgumentCaptor(ArgumentCaptor.forClass(T::class.java), T::class)
 }
 
+/**
+ * Creates a [KArgumentCaptor] for given type, taking in a lambda to allow fast verification.
+ */
+inline fun <reified T : Any> argumentCaptor(f: KArgumentCaptor<T>.() -> Unit): KArgumentCaptor<T> {
+    return argumentCaptor<T>().apply(f)
+}
+
 /**
  * Creates a [KArgumentCaptor] for given nullable type.
  */
@@ -43,6 +50,13 @@ inline fun <reified T : Any> nullableArgumentCaptor(): KArgumentCaptor<T?> {
     return KArgumentCaptor(ArgumentCaptor.forClass(T::class.java), T::class)
 }
 
+/**
+ * Creates a [KArgumentCaptor] for given nullable type, taking in a lambda to allow fast verification.
+ */
+inline fun <reified T : Any> nullableArgumentCaptor(f: KArgumentCaptor<T?>.() -> Unit): KArgumentCaptor<T?> {
+    return nullableArgumentCaptor<T>().apply(f)
+}
+
 /**
  * Alias for [ArgumentCaptor.capture].
  */
diff --git a/tests/src/test/kotlin/test/ArgumentCaptorTest.kt b/tests/src/test/kotlin/test/ArgumentCaptorTest.kt
index 551cc6b9..b79227b2 100644
--- a/tests/src/test/kotlin/test/ArgumentCaptorTest.kt
+++ b/tests/src/test/kotlin/test/ArgumentCaptorTest.kt
@@ -1,6 +1,7 @@
 package test
 
 import com.nhaarman.expect.expect
+import com.nhaarman.expect.expectErrorWithMessage
 import com.nhaarman.mockitokotlin2.*
 import org.junit.Test
 import java.util.*
@@ -117,4 +118,36 @@ class ArgumentCaptorTest : TestBase() {
             expect(secondValue).toBe(2)
         }
     }
+
+    @Test
+    fun argumentCaptor_withSingleValue_lambda() {
+        /* Given */
+        val date: Date = mock()
+
+        /* When */
+        date.time = 5L
+
+        /* Then */
+        argumentCaptor<Long> {
+            verify(date).time = capture()
+            expect(lastValue).toBe(5L)
+        }
+    }
+
+    @Test
+    fun argumentCaptor_withSingleValue_lambda_properlyFails() {
+        /* Given */
+        val date: Date = mock()
+
+        /* When */
+        date.time = 5L
+
+        /* Then */
+        expectErrorWithMessage("Expected: 3 but was: 5") on {
+            argumentCaptor<Long> {
+                verify(date).time = capture()
+                expect(lastValue).toBe(3L)
+            }
+        }
+    }
 }