Merge pull request #2365 from element-hq/feature/bma/testRoomList

Add Unit tests on MessagesView
2024-02-08 15:02:47 +01:00 · 2024-02-08 15:02:47 +01:00 · 5527c9634b
commit 5527c9634b
parent 89bcd378c5 e63e5bd739
9 changed files with 394 additions and 51 deletions
--- a/features/messages/impl/build.gradle.kts
+++ b/features/messages/impl/build.gradle.kts
@ -62,6 +62,7 @@ dependencies {
    implementation(projects.libraries.voicerecorder.api)
    implementation(projects.libraries.mediaplayer.api)
    implementation(projects.libraries.uiUtils)
+    implementation(projects.libraries.testtags)
    implementation(projects.features.networkmonitor.api)
    implementation(projects.services.analytics.api)
    implementation(libs.coil.compose)
--- a/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/MessagesStateProvider.kt
+++ b/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/MessagesStateProvider.kt
@ -17,16 +17,21 @@
 package io.element.android.features.messages.impl

 import androidx.compose.ui.tooling.preview.PreviewParameterProvider
+import io.element.android.features.messages.impl.actionlist.ActionListState
 import io.element.android.features.messages.impl.actionlist.anActionListState
 import io.element.android.features.messages.impl.messagecomposer.AttachmentsState
+import io.element.android.features.messages.impl.messagecomposer.MessageComposerState
 import io.element.android.features.messages.impl.messagecomposer.aMessageComposerState
 import io.element.android.features.messages.impl.timeline.aTimelineItemList
 import io.element.android.features.messages.impl.timeline.aTimelineState
+import io.element.android.features.messages.impl.timeline.components.customreaction.CustomReactionEvents
 import io.element.android.features.messages.impl.timeline.components.customreaction.CustomReactionState
+import io.element.android.features.messages.impl.timeline.components.reactionsummary.ReactionSummaryEvents
 import io.element.android.features.messages.impl.timeline.components.reactionsummary.ReactionSummaryState
 import io.element.android.features.messages.impl.timeline.components.receipt.bottomsheet.ReadReceiptBottomSheetState
 import io.element.android.features.messages.impl.timeline.components.retrysendmenu.RetrySendMenuState
 import io.element.android.features.messages.impl.timeline.model.event.aTimelineItemTextContent
+import io.element.android.features.messages.impl.voicemessages.composer.VoiceMessageComposerState
 import io.element.android.features.messages.impl.voicemessages.composer.aVoiceMessageComposerState
 import io.element.android.features.messages.impl.voicemessages.composer.aVoiceMessagePreviewState
 import io.element.android.libraries.architecture.AsyncData
@ -42,60 +47,74 @@ open class MessagesStateProvider : PreviewParameterProvider<MessagesState> {
    override val values: Sequence<MessagesState>
        get() = sequenceOf(
            aMessagesState(),
-            aMessagesState().copy(hasNetworkConnection = false),
-            aMessagesState().copy(composerState = aMessageComposerState().copy(showAttachmentSourcePicker = true)),
-            aMessagesState().copy(userHasPermissionToSendMessage = false),
-            aMessagesState().copy(showReinvitePrompt = true),
-            aMessagesState().copy(
+            aMessagesState(hasNetworkConnection = false),
+            aMessagesState(composerState = aMessageComposerState(showAttachmentSourcePicker = true)),
+            aMessagesState(userHasPermissionToSendMessage = false),
+            aMessagesState(showReinvitePrompt = true),
+            aMessagesState(
                roomName = AsyncData.Uninitialized,
                roomAvatar = AsyncData.Uninitialized,
            ),
-            aMessagesState().copy(composerState = aMessageComposerState().copy(showTextFormatting = true)),
-            aMessagesState().copy(
+            aMessagesState(composerState = aMessageComposerState(showTextFormatting = true)),
+            aMessagesState(
                enableVoiceMessages = true,
                voiceMessageComposerState = aVoiceMessageComposerState(showPermissionRationaleDialog = true),
            ),
-            aMessagesState().copy(
-                composerState = aMessageComposerState().copy(
+            aMessagesState(
+                composerState = aMessageComposerState(
                    attachmentsState = AttachmentsState.Sending.Processing(persistentListOf())
                ),
            ),
-            aMessagesState().copy(
-                composerState = aMessageComposerState().copy(
+            aMessagesState(
+                composerState = aMessageComposerState(
                    attachmentsState = AttachmentsState.Sending.Uploading(0.33f)
                ),
            ),
-            aMessagesState().copy(
+            aMessagesState(
                callState = RoomCallState.ONGOING,
            ),
-            aMessagesState().copy(
+            aMessagesState(
                enableVoiceMessages = true,
                voiceMessageComposerState = aVoiceMessageComposerState(
                    voiceMessageState = aVoiceMessagePreviewState(),
                    showSendFailureDialog = true
                ),
            ),
-            aMessagesState().copy(
+            aMessagesState(
                callState = RoomCallState.DISABLED,
            ),
        )
 }

-fun aMessagesState() = MessagesState(
-    roomId = RoomId("!id:domain"),
-    roomName = AsyncData.Success("Room name"),
-    roomAvatar = AsyncData.Success(AvatarData("!id:domain", "Room name", size = AvatarSize.TimelineRoom)),
-    userHasPermissionToSendMessage = true,
-    userHasPermissionToRedactOwn = false,
-    userHasPermissionToRedactOther = false,
-    userHasPermissionToSendReaction = true,
-    composerState = aMessageComposerState().copy(
+fun aMessagesState(
+    roomName: AsyncData<String> = AsyncData.Success("Room name"),
+    roomAvatar: AsyncData<AvatarData> = AsyncData.Success(AvatarData("!id:domain", "Room name", size = AvatarSize.TimelineRoom)),
+    userHasPermissionToSendMessage: Boolean = true,
+    composerState: MessageComposerState = aMessageComposerState(
        richTextEditorState = RichTextEditorState("Hello", initialFocus = true),
        isFullScreen = false,
        mode = MessageComposerMode.Normal,
    ),
-    voiceMessageComposerState = aVoiceMessageComposerState(),
-    timelineState = aTimelineState().copy(
+    voiceMessageComposerState: VoiceMessageComposerState = aVoiceMessageComposerState(),
+    actionListState: ActionListState = anActionListState(),
+    customReactionState: CustomReactionState = aCustomReactionState(),
+    reactionSummaryState: ReactionSummaryState = aReactionSummaryState(),
+    hasNetworkConnection: Boolean = true,
+    showReinvitePrompt: Boolean = false,
+    enableVoiceMessages: Boolean = true,
+    callState: RoomCallState = RoomCallState.ENABLED,
+    eventSink: (MessagesEvents) -> Unit = {},
+) = MessagesState(
+    roomId = RoomId("!id:domain"),
+    roomName = roomName,
+    roomAvatar = roomAvatar,
+    userHasPermissionToSendMessage = userHasPermissionToSendMessage,
+    userHasPermissionToRedactOwn = false,
+    userHasPermissionToRedactOther = false,
+    userHasPermissionToSendReaction = true,
+    composerState = composerState,
+    voiceMessageComposerState = voiceMessageComposerState,
+    timelineState = aTimelineState(
        timelineItems = aTimelineItemList(aTimelineItemTextContent()),
    ),
    retrySendMenuState = RetrySendMenuState(
@ -106,23 +125,32 @@ fun aMessagesState() = MessagesState(
        selectedEvent = null,
        eventSink = {},
    ),
-    actionListState = anActionListState(),
-    customReactionState = CustomReactionState(
-        target = CustomReactionState.Target.None,
-        eventSink = {},
-        selectedEmoji = persistentSetOf(),
-    ),
-    reactionSummaryState = ReactionSummaryState(
-        target = null,
-        eventSink = {},
-    ),
-    hasNetworkConnection = true,
+    actionListState = actionListState,
+    customReactionState = customReactionState,
+    reactionSummaryState = reactionSummaryState,
+    hasNetworkConnection = hasNetworkConnection,
    snackbarMessage = null,
    inviteProgress = AsyncData.Uninitialized,
-    showReinvitePrompt = false,
+    showReinvitePrompt = showReinvitePrompt,
    enableTextFormatting = true,
-    enableVoiceMessages = true,
-    callState = RoomCallState.ENABLED,
+    enableVoiceMessages = enableVoiceMessages,
+    callState = callState,
    appName = "Element",
-    eventSink = {}
+    eventSink = eventSink,
+)
+
+fun aReactionSummaryState(
+    target: ReactionSummaryState.Summary? = null,
+    eventSink: (ReactionSummaryEvents) -> Unit = {}
+) = ReactionSummaryState(
+    target = target,
+    eventSink = eventSink,
+)
+
+fun aCustomReactionState(
+    eventSink: (CustomReactionEvents) -> Unit = {},
+) = CustomReactionState(
+    target = CustomReactionState.Target.None,
+    selectedEmoji = persistentSetOf(),
+    eventSink = eventSink,
 )
--- a/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/actionlist/ActionListStateProvider.kt
+++ b/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/actionlist/ActionListStateProvider.kt
@ -122,9 +122,12 @@ open class ActionListStateProvider : PreviewParameterProvider<ActionListState> {
        }
 }

-fun anActionListState() = ActionListState(
-    target = ActionListState.Target.None,
-    eventSink = {}
+fun anActionListState(
+    target: ActionListState.Target = ActionListState.Target.None,
+    eventSink: (ActionListEvents) -> Unit = {},
+) = ActionListState(
+    target = target,
+    eventSink = eventSink
 )

 fun aTimelineItemActionList(): ImmutableList<TimelineItemAction> {
--- a/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/messagecomposer/MessageComposerStateProvider.kt
+++ b/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/messagecomposer/MessageComposerStateProvider.kt
@ -32,7 +32,7 @@ open class MessageComposerStateProvider : PreviewParameterProvider<MessageCompos
 }

 fun aMessageComposerState(
-    composerState: RichTextEditorState = RichTextEditorState(""),
+    richTextEditorState: RichTextEditorState = RichTextEditorState(""),
    isFullScreen: Boolean = false,
    mode: MessageComposerMode = MessageComposerMode.Normal,
    showTextFormatting: Boolean = false,
@ -42,7 +42,7 @@ fun aMessageComposerState(
    attachmentsState: AttachmentsState = AttachmentsState.None,
    memberSuggestions: ImmutableList<MentionSuggestion> = persistentListOf(),
 ) = MessageComposerState(
-    richTextEditorState = composerState,
+    richTextEditorState = richTextEditorState,
    isFullScreen = isFullScreen,
    mode = mode,
    showTextFormatting = showTextFormatting,
--- a/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/timeline/components/MessageEventBubble.kt
+++ b/features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/timeline/components/MessageEventBubble.kt
@ -46,6 +46,8 @@ import io.element.android.libraries.designsystem.theme.components.Surface
 import io.element.android.libraries.designsystem.theme.components.Text
 import io.element.android.libraries.designsystem.theme.messageFromMeBackground
 import io.element.android.libraries.designsystem.theme.messageFromOtherBackground
+import io.element.android.libraries.testtags.TestTags
+import io.element.android.libraries.testtags.testTag

 private val BUBBLE_RADIUS = 12.dp
 internal val BUBBLE_INCOMING_OFFSET = 16.dp
@ -115,6 +117,7 @@ fun MessageEventBubble(
    ) {
        Surface(
            modifier = Modifier
+                .testTag(TestTags.messageBubble)
                .widthIn(min = 80.dp)
                .clip(bubbleShape)
                .combinedClickable(
--- a/features/messages/impl/src/test/kotlin/io/element/android/features/messages/impl/MessagesViewTest.kt
+++ b/features/messages/impl/src/test/kotlin/io/element/android/features/messages/impl/MessagesViewTest.kt
@ -0,0 +1,288 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.element.android.features.messages.impl
+
+import androidx.activity.ComponentActivity
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.ui.platform.LocalInspectionMode
+import androidx.compose.ui.test.junit4.AndroidComposeTestRule
+import androidx.compose.ui.test.junit4.createAndroidComposeRule
+import androidx.compose.ui.test.longClick
+import androidx.compose.ui.test.onAllNodesWithContentDescription
+import androidx.compose.ui.test.onAllNodesWithTag
+import androidx.compose.ui.test.onAllNodesWithText
+import androidx.compose.ui.test.onFirst
+import androidx.compose.ui.test.onNodeWithContentDescription
+import androidx.compose.ui.test.onNodeWithText
+import androidx.compose.ui.test.performClick
+import androidx.compose.ui.test.performTouchInput
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import io.element.android.features.messages.impl.actionlist.ActionListState
+import io.element.android.features.messages.impl.actionlist.anActionListState
+import io.element.android.features.messages.impl.actionlist.model.TimelineItemAction
+import io.element.android.features.messages.impl.attachments.Attachment
+import io.element.android.features.messages.impl.messagecomposer.aMessageComposerState
+import io.element.android.features.messages.impl.timeline.components.customreaction.CustomReactionEvents
+import io.element.android.features.messages.impl.timeline.components.reactionsummary.ReactionSummaryEvents
+import io.element.android.features.messages.impl.timeline.model.TimelineItem
+import io.element.android.libraries.matrix.api.core.UserId
+import io.element.android.libraries.testtags.TestTags
+import io.element.android.libraries.ui.strings.CommonStrings
+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.EnsureNeverCalledWithParamAndResult
+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 kotlinx.collections.immutable.ImmutableList
+import kotlinx.collections.immutable.persistentListOf
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TestRule
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class MessagesViewTest {
+    @get:Rule val rule = createAndroidComposeRule<ComponentActivity>()
+
+    @Test
+    fun `clicking on back invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        ensureCalledOnce { callback ->
+            rule.setMessagesView(
+                state = state,
+                onBackPressed = callback,
+            )
+            rule.pressBack()
+        }
+    }
+
+    @Test
+    fun `clicking on room name invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        ensureCalledOnce { callback ->
+            rule.setMessagesView(
+                state = state,
+                onRoomDetailsClicked = callback,
+            )
+            rule.onNodeWithText(state.roomName.dataOrNull().orEmpty()).performClick()
+        }
+    }
+
+    @Test
+    fun `clicking on join call invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        ensureCalledOnce { callback ->
+            rule.setMessagesView(
+                state = state,
+                onJoinCallClicked = callback,
+            )
+            val joinCallContentDescription = rule.activity.getString(CommonStrings.a11y_start_call)
+            rule.onNodeWithContentDescription(joinCallContentDescription).performClick()
+        }
+    }
+
+    @Test
+    fun `clicking on an Event invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        val timelineItem = state.timelineState.timelineItems.first()
+        val callback = EnsureCalledOnceWithParam(
+            expectedParam = timelineItem,
+            result = true,
+        )
+        rule.setMessagesView(
+            state = state,
+            onEventClicked = callback,
+        )
+        // Cannot perform click on "Text", it's not detected. Use tag instead
+        rule.onAllNodesWithTag(TestTags.messageBubble.value).onFirst().performClick()
+        callback.assertSuccess()
+    }
+
+    @Test
+    fun `clicking on send location invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            composerState = aMessageComposerState(
+                showAttachmentSourcePicker = true
+            ),
+            eventSink = eventsRecorder
+        )
+        ensureCalledOnce { callback ->
+            rule.setMessagesView(
+                state = state,
+                onSendLocationClicked = callback,
+            )
+            rule.clickOn(R.string.screen_room_attachment_source_location)
+        }
+    }
+
+    @Test
+    fun `clicking on create poll invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            composerState = aMessageComposerState(
+                showAttachmentSourcePicker = true
+            ),
+            eventSink = eventsRecorder
+        )
+        ensureCalledOnce { callback ->
+            rule.setMessagesView(
+                state = state,
+                onCreatePollClicked = callback,
+            )
+            // Then click on the poll action
+            rule.clickOn(R.string.screen_room_attachment_source_poll)
+        }
+    }
+
+    @Test
+    fun `clicking on the sender of an Event invoke expected callback`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>(expectEvents = false)
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        val timelineItem = state.timelineState.timelineItems.first()
+        ensureCalledOnceWithParam(
+            param = (timelineItem as TimelineItem.Event).senderId
+        ) { callback ->
+            rule.setMessagesView(
+                state = state,
+                onUserDataClicked = callback,
+            )
+            val senderName = (timelineItem as? TimelineItem.Event)?.senderDisplayName.orEmpty()
+            rule.onNodeWithText(senderName).performClick()
+        }
+    }
+
+    @Test
+    fun `selecting a action on a message emits the expected Event`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>()
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        val timelineItem = state.timelineState.timelineItems.first() as TimelineItem.Event
+        val stateWithMessageAction = state.copy(
+            actionListState = anActionListState(
+                target = ActionListState.Target.Success(
+                    event = timelineItem,
+                    displayEmojiReactions = true,
+                    actions = persistentListOf(TimelineItemAction.Edit),
+                )
+            ),
+        )
+        rule.setMessagesView(
+            state = stateWithMessageAction,
+        )
+        rule.clickOn(CommonStrings.action_edit)
+        // Give time for the close animation to complete
+        rule.mainClock.advanceTimeBy(milliseconds = 1_000)
+        eventsRecorder.assertSingle(MessagesEvents.HandleAction(TimelineItemAction.Edit, timelineItem))
+    }
+
+    @Test
+    fun `clicking on a reaction emits the expected Event`() {
+        val eventsRecorder = EventsRecorder<MessagesEvents>()
+        val state = aMessagesState(
+            eventSink = eventsRecorder
+        )
+        val timelineItem = state.timelineState.timelineItems.first() as TimelineItem.Event
+        rule.setMessagesView(
+            state = state,
+        )
+        rule.onAllNodesWithText("👍️").onFirst().performClick()
+        eventsRecorder.assertSingle(MessagesEvents.ToggleReaction("👍️", timelineItem.eventId!!))
+    }
+
+    @Test
+    fun `long clicking on a reaction emits the expected Event`() {
+        val eventsRecorder = EventsRecorder<ReactionSummaryEvents>()
+        val state = aMessagesState(
+            reactionSummaryState = aReactionSummaryState(
+                target = null,
+                eventSink = eventsRecorder,
+            ),
+        )
+        val timelineItem = state.timelineState.timelineItems.first() as TimelineItem.Event
+        rule.setMessagesView(
+            state = state,
+        )
+        rule.onAllNodesWithText("👍️").onFirst().performTouchInput { longClick() }
+        eventsRecorder.assertSingle(ReactionSummaryEvents.ShowReactionSummary(timelineItem.eventId!!, timelineItem.reactionsState.reactions, "👍️"))
+    }
+
+    @Test
+    fun `clicking on more reaction emits the expected Event`() {
+        val eventsRecorder = EventsRecorder<CustomReactionEvents>()
+        val state = aMessagesState(
+            customReactionState = aCustomReactionState(
+                eventSink = eventsRecorder,
+            ),
+        )
+        val timelineItem = state.timelineState.timelineItems.first() as TimelineItem.Event
+        rule.setMessagesView(
+            state = state,
+        )
+        val moreReactionContentDescription = rule.activity.getString(R.string.screen_room_timeline_add_reaction)
+        rule.onAllNodesWithContentDescription(moreReactionContentDescription).onFirst().performClick()
+        eventsRecorder.assertSingle(CustomReactionEvents.ShowCustomReactionSheet(timelineItem))
+    }
+}
+
+private fun <R : TestRule> AndroidComposeTestRule<R, ComponentActivity>.setMessagesView(
+    state: MessagesState,
+    onBackPressed: () -> Unit = EnsureNeverCalled(),
+    onRoomDetailsClicked: () -> Unit = EnsureNeverCalled(),
+    onEventClicked: (event: TimelineItem.Event) -> Boolean = EnsureNeverCalledWithParamAndResult(),
+    onUserDataClicked: (UserId) -> Unit = EnsureNeverCalledWithParam(),
+    onPreviewAttachments: (ImmutableList<Attachment>) -> Unit = EnsureNeverCalledWithParam(),
+    onSendLocationClicked: () -> Unit = EnsureNeverCalled(),
+    onCreatePollClicked: () -> Unit = EnsureNeverCalled(),
+    onJoinCallClicked: () -> Unit = EnsureNeverCalled(),
+) {
+    setContent {
+        // Cannot use the RichTextEditor, so simulate a LocalInspectionMode
+        CompositionLocalProvider(LocalInspectionMode provides true) {
+            MessagesView(
+                state = state,
+                onBackPressed = onBackPressed,
+                onRoomDetailsClicked = onRoomDetailsClicked,
+                onEventClicked = onEventClicked,
+                onUserDataClicked = onUserDataClicked,
+                onPreviewAttachments = onPreviewAttachments,
+                onSendLocationClicked = onSendLocationClicked,
+                onCreatePollClicked = onCreatePollClicked,
+                onJoinCallClicked = onJoinCallClicked,
+            )
+        }
+    }
+}
--- a/libraries/testtags/src/main/kotlin/io/element/android/libraries/testtags/TestTags.kt
+++ b/libraries/testtags/src/main/kotlin/io/element/android/libraries/testtags/TestTags.kt
@ -58,6 +58,11 @@ object TestTags {
     */
    val richTextEditor = TestTag("rich_text_editor")

+    /**
+     * Message bubble.
+     */
+    val messageBubble = TestTag("message_bubble")
+
    /**
     * Dialogs.
     */
--- 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()
 }
--- a/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureNeverCalled.kt
+++ b/tests/testutils/src/main/kotlin/io/element/android/tests/testutils/EnsureNeverCalled.kt
@ -28,6 +28,12 @@ class EnsureNeverCalledWithParam<T> : (T) -> Unit {
    }
 }

+class EnsureNeverCalledWithParamAndResult<T, R> : (T) -> R {
+    override fun invoke(p1: T): R {
+        throw AssertionError("Should not be called and is called with $p1")
+    }
+}
+
 class EnsureNeverCalledWithTwoParams<T, U> : (T, U) -> Unit {
    override fun invoke(p1: T, p2: U) {
        throw AssertionError("Should not be called and is called with $p1 and $p2")