From 0628ceabc1f1196b613abef2fd2c6c84e184953a Mon Sep 17 00:00:00 2001
From: Benoit Marty <benoit@matrix.org>
Date: Mon, 8 Jan 2024 10:07:48 +0100
Subject: [PATCH] Iterate on EnsureCalledOnce devX

---
 .../features/logout/impl/LogoutViewTest.kt    | 10 +++++-----
 .../tests/testutils/EnsureCalledOnce.kt       | 20 ++++++++++---------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutViewTest.kt b/features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutViewTest.kt
index f0abf45234..4958afabe7 100644
--- a/features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutViewTest.kt
+++ b/features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutViewTest.kt
@@ -21,12 +21,12 @@ import androidx.compose.ui.test.junit4.createAndroidComposeRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import io.element.android.libraries.architecture.AsyncAction
 import io.element.android.libraries.ui.strings.CommonStrings
-import io.element.android.tests.testutils.EnsureCalledOnce
-import io.element.android.tests.testutils.EnsureCalledOnceWithParam
 import io.element.android.tests.testutils.EnsureNeverCalled
 import io.element.android.tests.testutils.EnsureNeverCalledWithParam
 import io.element.android.tests.testutils.EventsRecorder
 import io.element.android.tests.testutils.clickOn
+import io.element.android.tests.testutils.ensureCalledOnce
+import io.element.android.tests.testutils.ensureCalledOnceWithParam
 import io.element.android.tests.testutils.pressBack
 import org.junit.Rule
 import org.junit.Test
@@ -57,7 +57,7 @@ class LogoutViewTest {
     @Test
     fun `clicking on back invoke back callback`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        EnsureCalledOnce().run { callback ->
+        ensureCalledOnce { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
@@ -112,7 +112,7 @@ class LogoutViewTest {
     fun `success logout invoke onSuccessLogout`() {
         val data = "data"
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        EnsureCalledOnceWithParam<String?>(data).run { callback ->
+        ensureCalledOnceWithParam<String?>(data) { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
@@ -130,7 +130,7 @@ class LogoutViewTest {
     @Test
     fun `last session setting button invoke onChangeRecoveryKeyClicked`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        EnsureCalledOnce().run { callback ->
+        ensureCalledOnce { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
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 a29c320a00..ddc8e22a5c 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
@@ -27,11 +27,12 @@ class EnsureCalledOnce : () -> Unit {
             throw AssertionError("Expected to be called once, but was called $counter times")
         }
     }
+}
 
-    fun run(block: (callback: EnsureCalledOnce) -> Unit) {
-        block(this)
-        assertSuccess()
-    }
+fun ensureCalledOnce(block: (callback: EnsureCalledOnce) -> Unit) {
+    val callback = EnsureCalledOnce()
+    block(callback)
+    callback.assertSuccess()
 }
 
 class EnsureCalledOnceWithParam<T>(
@@ -50,9 +51,10 @@ class EnsureCalledOnceWithParam<T>(
             throw AssertionError("Expected to be called once, but was called $counter times")
         }
     }
-
-    fun run(block: (callback: EnsureCalledOnceWithParam<T>) -> Unit) {
-        block(this)
-        assertSuccess()
-    }
+}
+
+fun <T> ensureCalledOnceWithParam(param: T, block: (callback: EnsureCalledOnceWithParam<T>) -> Unit) {
+    val callback = EnsureCalledOnceWithParam(param)
+    block(callback)
+    callback.assertSuccess()
 }