Test DefaultWorkManagerScheduler

libraries/session-storage/test/src/main/kotlin/io/element/android/libraries/sessionstorage/test/InMemorySessionStore.kt

@@ -81,8 +81,6 @@ class InMemorySessionStore(
     }
 
     override suspend fun removeSession(sessionId: String) {
-        val currentList = sessionDataListFlow.value.toMutableList()
+        sessionDataListFlow.value = sessionDataListFlow.value.filter { it.userId != sessionId }
-        currentList.removeAll { it.userId == sessionId }
-        sessionDataListFlow.value = currentList
     }
 }

libraries/workmanager/impl/build.gradle.kts

@@ -1,4 +1,5 @@
 import extension.setupDependencyInjection
+import extension.testCommonDependencies
 
 /*
  * Copyright (c) 2025 Element Creations Ltd.
@@ -22,4 +23,7 @@ dependencies {
     implementation(projects.libraries.core)
     implementation(projects.libraries.matrix.api)
     implementation(projects.libraries.di)
+
+    testCommonDependencies(libs, false)
+    testImplementation(projects.libraries.sessionStorage.test)
 }

libraries/workmanager/impl/src/main/kotlin/io/element/android/libraries/workmanager/impl/DefaultWorkManagerScheduler.kt

@@ -8,23 +8,45 @@
 
 package io.element.android.libraries.workmanager.impl
 
-import android.content.Context
 import androidx.work.WorkManager
 import dev.zacsweers.metro.AppScope
 import dev.zacsweers.metro.ContributesBinding
-import io.element.android.libraries.di.annotations.ApplicationContext
+import io.element.android.libraries.core.coroutine.withPreviousValue
+import io.element.android.libraries.di.annotations.AppCoroutineScope
 import io.element.android.libraries.matrix.api.core.SessionId
+import io.element.android.libraries.sessionstorage.api.SessionStore
 import io.element.android.libraries.workmanager.api.WorkManagerRequest
 import io.element.android.libraries.workmanager.api.WorkManagerRequestType
 import io.element.android.libraries.workmanager.api.WorkManagerScheduler
 import io.element.android.libraries.workmanager.api.workManagerTag
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import timber.log.Timber
 
 @ContributesBinding(AppScope::class)
 class DefaultWorkManagerScheduler(
-    @ApplicationContext private val context: Context,
+    lazyWorkManager: Lazy<WorkManager>,
+    @AppCoroutineScope private val appCoroutineScope: CoroutineScope,
+    sessionStore: SessionStore,
 ) : WorkManagerScheduler {
-    private val workManager by lazy { WorkManager.getInstance(context) }
+    private val workManager by lazyWorkManager
+
+    init {
+        // Observe session removals to cancel associated work automatically
+        sessionStore.sessionsFlow()
+            .map { sessions -> sessions.map { SessionId(it.userId) } }
+            .withPreviousValue()
+            .map { (prev, new) -> prev.orEmpty() - new.toSet() }
+            .onEach { removedSessions ->
+                for (sessionId in removedSessions) {
+                    Timber.d("Session removed for userId: $sessionId, cancelling associated workmanager requests")
+                    cancel(sessionId)
+                }
+            }
+            .launchIn(appCoroutineScope)
+    }
 
     override fun submit(workManagerRequest: WorkManagerRequest) {
         workManagerRequest.build().fold(

libraries/workmanager/impl/src/main/kotlin/io/element/android/libraries/workmanager/impl/WorkManagerModule.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2025 Element Creations 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.workmanager.impl
+
+import android.content.Context
+import androidx.work.WorkManager
+import dev.zacsweers.metro.AppScope
+import dev.zacsweers.metro.BindingContainer
+import dev.zacsweers.metro.ContributesTo
+import dev.zacsweers.metro.Provides
+import dev.zacsweers.metro.SingleIn
+import io.element.android.libraries.di.annotations.ApplicationContext
+
+@BindingContainer
+@ContributesTo(AppScope::class)
+object WorkManagerModule {
+    @Provides
+    @SingleIn(AppScope::class)
+    fun providesWorkManager(
+        @ApplicationContext context: Context,
+    ): WorkManager {
+        return WorkManager.getInstance(context)
+    }
+}

libraries/workmanager/impl/src/test/kotlin/io/element/android/libraries/workmanager/impl/DefaultWorkManagerSchedulerTest.kt

@@ -0,0 +1,120 @@
+/*
+ * Copyright (c) 2025 Element Creations 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.workmanager.impl
+
+import androidx.work.WorkManager
+import androidx.work.WorkRequest
+import com.google.common.truth.Truth.assertThat
+import io.element.android.libraries.matrix.api.core.SessionId
+import io.element.android.libraries.sessionstorage.test.InMemorySessionStore
+import io.element.android.libraries.sessionstorage.test.aSessionData
+import io.element.android.libraries.workmanager.api.WorkManagerRequest
+import io.element.android.libraries.workmanager.api.WorkManagerRequestType
+import io.element.android.libraries.workmanager.api.workManagerTag
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.spyk
+import io.mockk.verify
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class DefaultWorkManagerSchedulerTest {
+    @Test
+    fun `starts observing sessions on init to remove work for logged out sessions`() = runTest {
+        val sessionId = "@session1:matrix.org"
+        val sessionStore = InMemorySessionStore(initialList = listOf(aSessionData(sessionId = sessionId)))
+
+        val workManager = spyk<WorkManager>()
+
+        DefaultWorkManagerScheduler(
+            lazyWorkManager = lazy { workManager },
+            appCoroutineScope = backgroundScope,
+            sessionStore = sessionStore,
+        )
+
+        // We have a single initial session
+        assertThat(sessionStore.numberOfSessions()).isEqualTo(1)
+
+        runCurrent()
+
+        // We remove the session
+        sessionStore.removeSession(sessionId)
+
+        runCurrent()
+
+        // The session is now gone and work associated with the session is cancelled
+        assertThat(sessionStore.numberOfSessions()).isEqualTo(0)
+        verify { workManager.cancelAllWorkByTag("notifications-$sessionId") }
+    }
+
+    @Test
+    fun `submit builds the request and enqueues it`() = runTest {
+        val sessionStore = InMemorySessionStore()
+        val workManager = spyk<WorkManager>()
+
+        val scheduler = DefaultWorkManagerScheduler(
+            lazyWorkManager = lazy { workManager },
+            appCoroutineScope = backgroundScope,
+            sessionStore = sessionStore,
+        )
+
+        scheduler.submit(FakeWorkManagerRequest())
+
+        verify { workManager.enqueue(any<List<WorkRequest>>()) }
+    }
+
+    @Test
+    fun `submit won't do anything if building the work request fails`() = runTest {
+        val sessionStore = InMemorySessionStore()
+        val workManager = spyk<WorkManager>()
+
+        val scheduler = DefaultWorkManagerScheduler(
+            lazyWorkManager = lazy { workManager },
+            appCoroutineScope = backgroundScope,
+            sessionStore = sessionStore,
+        )
+
+        scheduler.submit(FakeWorkManagerRequest(result = Result.failure(IllegalStateException("Test error"))))
+
+        verify(exactly = 0) { workManager.enqueue(any<List<WorkRequest>>()) }
+    }
+
+    @Test
+    fun `cancel will cancel all pending work for a session id`() = runTest {
+        val sessionStore = InMemorySessionStore()
+        val workManager = spyk<WorkManager>()
+
+        val scheduler = DefaultWorkManagerScheduler(
+            lazyWorkManager = lazy { workManager },
+            appCoroutineScope = backgroundScope,
+            sessionStore = sessionStore,
+        )
+
+        val sessionId = SessionId("@alice:matrix.org")
+        val tagToRemove = workManagerTag(sessionId, WorkManagerRequestType.NOTIFICATION_SYNC)
+        val mockSessionA = mockk<WorkRequest> {
+            every { tags } returns setOf(tagToRemove)
+        }
+        scheduler.submit(FakeWorkManagerRequest(result = Result.success(listOf(mockSessionA))))
+
+        scheduler.cancel(sessionId)
+
+        verify { workManager.cancelAllWorkByTag(tagToRemove) }
+    }
+}
+
+private class FakeWorkManagerRequest(
+    private val result: Result<List<WorkRequest>> = Result.success(listOf()),
+) : WorkManagerRequest {
+    override fun build(): Result<List<WorkRequest>> {
+        return result
+    }
+}