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,7 +57,7 @@ class LogoutViewTest {
     @Test
     fun `clicking on back invoke back callback`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnce()
+        EnsureCalledOnce().run { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
@@ -69,7 +69,7 @@ class LogoutViewTest {
                 )
             }
             rule.pressBack()
-        callback.assertSuccess()
+        }
     }
 
     @Test
@@ -112,7 +112,7 @@ class LogoutViewTest {
     fun `success logout invoke onSuccessLogout`() {
         val data = "data"
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnceWithParam<String?>(data)
+        EnsureCalledOnceWithParam<String?>(data).run { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
@@ -124,13 +124,13 @@ class LogoutViewTest {
                     onSuccessLogout = callback,
                 )
             }
-        callback.assertSuccess()
+        }
     }
 
     @Test
     fun `last session setting button invoke onChangeRecoveryKeyClicked`() {
         val eventsRecorder = EventsRecorder<LogoutEvents>(expectEvents = false)
-        val callback = EnsureCalledOnce()
+        EnsureCalledOnce().run { callback ->
             rule.setContent {
                 LogoutView(
                     aLogoutState(
@@ -143,6 +143,6 @@ class LogoutViewTest {
                 )
             }
             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()
+    }
 }