Add run method to EnsureCalledOnce and EnsureCalledOnceWithParam to ensure that assertSuccess is always called.

features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutViewTest.kt

@@ -57,19 +57,19 @@ class LogoutViewTest {
     @Test
     fun `clicking on back invoke back callback`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnce()
-        rule.setContent {
-            LogoutView(
-                aLogoutState(
-                    eventSink = eventsRecorder
-                ),
-                onChangeRecoveryKeyClicked = EnsureNeverCalled(),
-                onBackClicked = callback,
-                onSuccessLogout = EnsureNeverCalledWithParam(),
-            )
+        EnsureCalledOnce().run { callback ->
+            rule.setContent {
+                LogoutView(
+                    aLogoutState(
+                        eventSink = eventsRecorder
+                    ),
+                    onChangeRecoveryKeyClicked = EnsureNeverCalled(),
+                    onBackClicked = callback,
+                    onSuccessLogout = EnsureNeverCalledWithParam(),
+                )
+            }
+            rule.pressBack()
         }
-        rule.pressBack()
-        callback.assertSuccess()
     }
 
     @Test
@@ -112,37 +112,37 @@ class LogoutViewTest {
     fun `success logout invoke onSuccessLogout`() {
         val data = "data"
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnceWithParam<String?>(data)
-        rule.setContent {
-            LogoutView(
-                aLogoutState(
-                    logoutAction = AsyncAction.Success(data),
-                    eventSink = eventsRecorder
-                ),
-                onChangeRecoveryKeyClicked = EnsureNeverCalled(),
-                onBackClicked = EnsureNeverCalled(),
-                onSuccessLogout = callback,
-            )
+        EnsureCalledOnceWithParam<String?>(data).run { callback ->
+            rule.setContent {
+                LogoutView(
+                    aLogoutState(
+                        logoutAction = AsyncAction.Success(data),
+                        eventSink = eventsRecorder
+                    ),
+                    onChangeRecoveryKeyClicked = EnsureNeverCalled(),
+                    onBackClicked = EnsureNeverCalled(),
+                    onSuccessLogout = callback,
+                )
+            }
         }
-        callback.assertSuccess()
     }
 
     @Test
     fun `last session setting button invoke onChangeRecoveryKeyClicked`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnce()
-        rule.setContent {
-            LogoutView(
-                aLogoutState(
-                    isLastSession = true,
-                    eventSink = eventsRecorder
-                ),
-                onChangeRecoveryKeyClicked = callback,
-                onBackClicked = EnsureNeverCalled(),
-                onSuccessLogout = EnsureNeverCalledWithParam(),
-            )
+        EnsureCalledOnce().run { callback ->
+            rule.setContent {
+                LogoutView(
+                    aLogoutState(
+                        isLastSession = true,
+                        eventSink = eventsRecorder
+                    ),
+                    onChangeRecoveryKeyClicked = callback,
+                    onBackClicked = EnsureNeverCalled(),
+                    onSuccessLogout = EnsureNeverCalledWithParam(),
+                )
+            }
+            rule.clickOn(CommonStrings.common_settings)
         }
-        rule.clickOn(CommonStrings.common_settings)
-        callback.assertSuccess()
     }
 }

tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureCalledOnce.kt

@@ -27,6 +27,11 @@ class EnsureCalledOnce : () -> Unit {
             throw AssertionError("Expected to be called once, but was called $counter times")
         }
     }
+
+    fun run(block: (callback: EnsureCalledOnce) -> Unit) {
+        block(this)
+        assertSuccess()
+    }
 }
 
 class EnsureCalledOnceWithParam<T>(
@@ -45,4 +50,9 @@ 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()
+    }
 }