From 6c96b53d9feab3236fa6285774ac5af4520dd20f Mon Sep 17 00:00:00 2001
From: Benoit Marty <benoit@matrix.org>
Date: Wed, 7 Feb 2024 17:22:12 +0100
Subject: [PATCH] Add result type to EnsureCalledOnceWithParam

---
 .../tests/testutils/EnsureCalledOnce.kt       | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureCalledOnce.kt b/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureCalledOnce.kt
index ddc8e22a5c..1d5fed72dc 100644
--- a/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureCalledOnce.kt
+++ b/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureCalledOnce.kt
@@ -35,15 +35,17 @@ fun ensureCalledOnce(block: (callback: EnsureCalledOnce) -> Unit) {
     callback.assertSuccess()
 }
 
-class EnsureCalledOnceWithParam<T>(
-    private val expectedParam: T
-) : (T) -> Unit {
+class EnsureCalledOnceWithParam<T, R>(
+    private val expectedParam: T,
+    private val result: R,
+) : (T) -> R {
     private var counter = 0
-    override fun invoke(p1: T) {
+    override fun invoke(p1: T): R {
         if (p1 != expectedParam) {
             throw AssertionError("Expected to be called with $expectedParam, but was called with $p1")
         }
         counter++
+        return result
     }
 
     fun assertSuccess() {
@@ -53,8 +55,15 @@ class EnsureCalledOnceWithParam<T>(
     }
 }
 
-fun <T> ensureCalledOnceWithParam(param: T, block: (callback: EnsureCalledOnceWithParam<T>) -> Unit) {
-    val callback = EnsureCalledOnceWithParam(param)
+/**
+ * Shortcut for [<T, R> ensureCalledOnceWithParam] with Unit result
+ */
+fun <T> ensureCalledOnceWithParam(param: T, block: (callback: EnsureCalledOnceWithParam<T, Unit>) -> Unit) {
+    ensureCalledOnceWithParam(param, block, Unit)
+}
+
+fun <T, R> ensureCalledOnceWithParam(param: T, block: (callback: EnsureCalledOnceWithParam<T, R>) -> R, result: R) {
+    val callback = EnsureCalledOnceWithParam(param, result)
     block(callback)
     callback.assertSuccess()
 }