From 959cb4f0b3ec66181c136f103ca3678d3a9486c8 Mon Sep 17 00:00:00 2001
From: Benoit Marty <benoit@matrix.org>
Date: Wed, 10 Sep 2025 18:18:16 +0200
Subject: [PATCH] Add test on SpaceRoomCache and fix implementation

---
 libraries/matrix/impl/build.gradle.kts        |  1 +
 .../matrix/impl/spaces/SpaceRoomCache.kt      | 25 +++++-----
 .../matrix/impl/spaces/SpaceRoomCacheTest.kt  | 46 +++++++++++++++++++
 3 files changed, 60 insertions(+), 12 deletions(-)
 create mode 100644 libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCacheTest.kt

diff --git a/libraries/matrix/impl/build.gradle.kts b/libraries/matrix/impl/build.gradle.kts
index a0fa346be2..757d7aa1fc 100644
--- a/libraries/matrix/impl/build.gradle.kts
+++ b/libraries/matrix/impl/build.gradle.kts
@@ -47,6 +47,7 @@ dependencies {
     testImplementation(projects.libraries.featureflag.test)
     testImplementation(projects.libraries.matrix.test)
     testImplementation(projects.libraries.preferences.test)
+    testImplementation(projects.libraries.previewutils)
     testImplementation(projects.libraries.sessionStorage.test)
     testImplementation(projects.services.analytics.test)
     testImplementation(projects.services.toolbox.test)
diff --git a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCache.kt b/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCache.kt
index 63fe020ecd..ce039b9471 100644
--- a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCache.kt
+++ b/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCache.kt
@@ -11,31 +11,32 @@ import io.element.android.libraries.matrix.api.core.RoomId
 import io.element.android.libraries.matrix.api.spaces.SpaceRoom
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
+import java.util.concurrent.ConcurrentHashMap
 
 /**
  * An in memory cache of space rooms.
  * Only caches Rooms with roomType [io.element.android.libraries.matrix.api.room.RoomType.Space].
  */
 class SpaceRoomCache {
-    private val inMemoryCache = MutableStateFlow<MutableMap<RoomId, SpaceRoom>>(LinkedHashMap())
+    private val inMemoryCache = ConcurrentHashMap<RoomId, MutableStateFlow<SpaceRoom?>>()
     private val mutex = Mutex()
 
     fun getSpaceRoomFlow(roomId: RoomId): Flow<SpaceRoom?> {
-        return inMemoryCache.map { it[roomId] }
+        return getMutableFlow(roomId).asStateFlow()
     }
 
     suspend fun update(spaceRooms: List<SpaceRoom>) = mutex.withLock {
-        inMemoryCache.update { cache ->
-            spaceRooms
-                .filter { it.isSpace }
-                .forEach { spaceRoom ->
-                    cache.put(spaceRoom.roomId, spaceRoom)
-                }
-            cache
-        }
+        spaceRooms
+            .filter { it.isSpace }
+            .forEach { spaceRoom ->
+                getMutableFlow(spaceRoom.roomId).value = spaceRoom
+            }
+    }
+
+    private fun getMutableFlow(roomId: RoomId): MutableStateFlow<SpaceRoom?> {
+        return inMemoryCache.getOrPut(roomId, { MutableStateFlow(null) })
     }
 }
diff --git a/libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCacheTest.kt b/libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCacheTest.kt
new file mode 100644
index 0000000000..0679ce1a33
--- /dev/null
+++ b/libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/spaces/SpaceRoomCacheTest.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2025 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.libraries.matrix.impl.spaces
+
+import app.cash.turbine.test
+import com.google.common.truth.Truth.assertThat
+import io.element.android.libraries.matrix.api.room.RoomType
+import io.element.android.libraries.matrix.test.A_ROOM_ID
+import io.element.android.libraries.matrix.test.A_ROOM_ID_2
+import io.element.android.libraries.previewutils.room.aSpaceRoom
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+
+class SpaceRoomCacheTest {
+    @Test
+    fun `getSpaceRoomFlow emits items`() = runTest {
+        val sut = SpaceRoomCache()
+        sut.getSpaceRoomFlow(A_ROOM_ID).test {
+            assertThat(awaitItem()).isNull()
+            val room = aSpaceRoom(
+                roomId = A_ROOM_ID,
+                roomType = RoomType.Room,
+            )
+            sut.update(listOf(room))
+            // Not a space, should not be cached
+            expectNoEvents()
+            val space = aSpaceRoom(
+                roomId = A_ROOM_ID,
+                roomType = RoomType.Space,
+            )
+            sut.update(listOf(space))
+            assertThat(awaitItem()).isEqualTo(space)
+            val spaceOther = aSpaceRoom(
+                roomId = A_ROOM_ID_2,
+                roomType = RoomType.Space,
+            )
+            sut.update(listOf(spaceOther))
+            expectNoEvents()
+        }
+    }
+}