Accept a stubbing lambda in mock()

mockito-kotlin/src/main/kotlin/com/nhaarman/mockito_kotlin/Mockito.kt

@@ -79,6 +79,18 @@ inline fun <reified T : Any> mock(defaultAnswer: Answer<Any>): T = Mockito.mock(
 inline fun <reified T : Any> mock(s: MockSettings): T = Mockito.mock(T::class.java, s)!!
 inline fun <reified T : Any> mock(s: String): T = Mockito.mock(T::class.java, s)!!
 
+inline fun <reified T : Any> mock(stubbing: KStubbing<T>.() -> Unit): T
+        = Mockito.mock(T::class.java)!!.apply { stubbing(KStubbing(this)) }
+
+class KStubbing<out T>(private val mock: T) {
+    fun <R> on(methodCall: R) = Mockito.`when`(methodCall)
+    fun <R> on(methodCall: T.() -> R) = Mockito.`when`(mock.methodCall())
+}
+
+infix fun <T> OngoingStubbing<T>.doReturn(t: T): OngoingStubbing<T> = thenReturn(t)
+fun <T> OngoingStubbing<T>.doReturn(t: T, vararg ts: T): OngoingStubbing<T> = thenReturn(t, *ts)
+inline infix fun <reified T> OngoingStubbing<T>.doReturn(ts: List<T>): OngoingStubbing<T> = thenReturn(ts[0], *ts.drop(1).toTypedArray())
+
 fun mockingDetails(toInspect: Any): MockingDetails = Mockito.mockingDetails(toInspect)!!
 fun never(): VerificationMode = Mockito.never()!!
 inline fun <reified T : Any> notNull(): T? = Mockito.notNull(T::class.java)

mockito-kotlin/src/test/kotlin/MockitoTest.kt

@@ -282,4 +282,60 @@ class MockitoTest {
             mock.go()
         }
     }
+
+    @Test
+    fun testMockStubbing_lambda() {
+        /* Given */
+        val mock = mock<Open>() {
+            on { stringResult() } doReturn "A"
+        }
+
+        /* When */
+        val result = mock.stringResult()
+
+        /* Then */
+        expect(result).toBe("A")
+    }
+
+    @Test
+    fun testMockStubbing_normalOverridesLambda() {
+        /* Given */
+        val mock = mock<Open>() {
+            on { stringResult() }.doReturn("A")
+        }
+        whenever(mock.stringResult()).thenReturn("B")
+
+        /* When */
+        val result = mock.stringResult()
+
+        /* Then */
+        expect(result).toBe("B")
+    }
+
+    @Test
+    fun testMockStubbing_methodCall() {
+        /* Given */
+        val mock = mock<Open>()
+        mock<Open> {
+            on(mock.stringResult()).doReturn("A")
+        }
+
+        /* When */
+        val result = mock.stringResult()
+
+        /* Then */
+        expect(result).toBe("A")
+    }
+
+    @Test
+    fun doReturn_withSingleItemList() {
+        /* Given */
+        val mock = mock<Open> {
+            on { stringResult() } doReturn listOf("a", "b")
+        }
+
+        /* Then */
+        expect(mock.stringResult()).toBe("a")
+        expect(mock.stringResult()).toBe("b")
+    }
 }