Add `doReturn().on { method() }` helper to KStubbing (#453) (#481)

Co-authored-by: Kate Corcoran <kate.corcoran@dfstudio.com>

2 files changed, 14 insertions, 0 deletions

diff --git a/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KStubbing.kt b/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KStubbing.kt
index aaf03b5..3ea1d88 100644
--- a/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KStubbing.kt
+++ b/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KStubbing.kt
@@ -30,6 +30,7 @@ import kotlinx.coroutines.runBlocking
 import org.mockito.Mockito
 import org.mockito.exceptions.misusing.NotAMockException
 import org.mockito.stubbing.OngoingStubbing
+import org.mockito.stubbing.Stubber
 import kotlin.reflect.KClass
 
 
@@ -85,4 +86,8 @@ class KStubbing<out T : Any>(val mock: T) {
     ): OngoingStubbing<R> {
         return runBlocking { Mockito.`when`(mock.m()) }
     }
+
+    fun Stubber.on(methodCall: T.() -> Unit) {
+        this.`when`(mock).methodCall()
+    }
 }
\ No newline at end of file
diff --git a/tests/src/test/kotlin/test/StubberTest.kt b/tests/src/test/kotlin/test/StubberTest.kt
index eb6e28c..81ac7f2 100644
--- a/tests/src/test/kotlin/test/StubberTest.kt
+++ b/tests/src/test/kotlin/test/StubberTest.kt
@@ -100,4 +100,13 @@ class StubberTest : TestBase() {
             mock.go()
         }
     }
+
+    @Test
+    fun testStubberOnBlockExtension() {
+        val mock = mock<Methods> {
+            doReturn("Test").on { stringResult() }
+        }
+
+        expect(mock.stringResult()).toBe("Test")
+    }
 }
\ No newline at end of file