Improve code again

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

@@ -0,0 +1,10 @@
+/*
+ * 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.push.impl.workmanager
+
+class DataForWorkManagerIsTooBig : Exception()

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

@@ -29,9 +29,9 @@ class SyncNotificationWorkManagerRequest(
     private val workerDataConverter: WorkerDataConverter,
     private val buildVersionSdkIntProvider: BuildVersionSdkIntProvider,
 ) : WorkManagerRequest {
-    override fun build(): List<Result<WorkRequest>> {
+    override fun build(): Result<List<WorkRequest>> {
         if (notificationEventRequests.isEmpty()) {
-            return listOf(Result.failure(InvalidParameterException("notificationEventRequests cannot be empty")))
+            return Result.failure(InvalidParameterException("notificationEventRequests cannot be empty"))
         }
         Timber.d("Scheduling ${notificationEventRequests.size} notification requests with WorkManager for $sessionId")
         return workerDataConverter.serialize(notificationEventRequests).map {

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

@@ -24,21 +24,22 @@ import timber.log.Timber
 class WorkerDataConverter(
     private val json: JsonProvider,
 ) {
-    fun serialize(notificationEventRequests: List<NotificationEventRequest>): List<Result<Data>> {
+    fun serialize(notificationEventRequests: List<NotificationEventRequest>): Result<List<Data>> {
         // First try to serialize all requests at once. In the vast majority of cases this will work.
         return serializeRequests(notificationEventRequests)
-            .fold(
-                onSuccess = {
-                    listOf(Result.success(it))
-                },
-                onFailure = {
-                    // Perform serialization on sublists, workDataOf may have failed because of size limit
+            .map { listOf(it) }
+            .recoverCatching {
+                if (it is DataForWorkManagerIsTooBig) {
+                    // Perform serialization on sublists, workDataOf have failed because of size limit
                     Timber.w(it, "Failed to serialize ${notificationEventRequests.size} notification requests, trying with chunks of $CHUNK_SIZE.")
-                    notificationEventRequests.chunked(CHUNK_SIZE).map { chunk ->
-                        serializeRequests(chunk)
+                    // TODO Do not split rooms
+                    notificationEventRequests.chunked(CHUNK_SIZE).mapNotNull { chunk ->
+                        serializeRequests(chunk).getOrNull()
                     }
-                },
-            )
+                } else {
+                    throw it
+                }
+            }
     }
 
     private fun serializeRequests(notificationEventRequests: List<NotificationEventRequest>): Result<Data> {
@@ -48,7 +49,11 @@ class WorkerDataConverter(
             }
             .mapCatchingExceptions { str ->
                 // Note: workDataOf can fail if the data is too large
-                workDataOf(REQUESTS_KEY to str)
+                try {
+                    workDataOf(REQUESTS_KEY to str)
+                } catch (_: IllegalStateException) {
+                    throw DataForWorkManagerIsTooBig()
+                }
             }
     }
 

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

@@ -33,8 +33,8 @@ class SyncNotificationWorkManagerRequestTest {
         )
 
         val result = request.build()
-        assertThat(result.first().isSuccess).isTrue()
-        result.first().getOrNull()!!.run {
+        assertThat(result.isSuccess).isTrue()
+        result.getOrNull()!!.first().run {
             assertThat(this).isInstanceOf(OneTimeWorkRequest::class.java)
             assertThat(workSpec.input.hasKeyWithValueOfType<String>("requests")).isTrue()
             // True in API 33+
@@ -52,8 +52,8 @@ class SyncNotificationWorkManagerRequestTest {
         )
 
         val result = request.build()
-        assertThat(result.first().isSuccess).isTrue()
-        result.first().getOrNull()!!.run {
+        assertThat(result.isSuccess).isTrue()
+        result.getOrNull()!!.first().run {
             assertThat(this).isInstanceOf(OneTimeWorkRequest::class.java)
             assertThat(workSpec.input.hasKeyWithValueOfType<String>("requests")).isTrue()
             // False before API 33
@@ -70,7 +70,7 @@ class SyncNotificationWorkManagerRequestTest {
         )
 
         val result = request.build()
-        assertThat(result.first().isFailure).isTrue()
+        assertThat(result.isFailure).isTrue()
     }
 
     @Test
@@ -81,7 +81,7 @@ class SyncNotificationWorkManagerRequestTest {
             workerDataConverter = WorkerDataConverter({ error("error during serialization") })
         )
         val result = request.build()
-        assertThat(result.first().isFailure).isTrue()
+        assertThat(result.isFailure).isTrue()
     }
 }
 

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

@@ -58,14 +58,14 @@ class WorkerDataConverterTest {
         }
         val sut = WorkerDataConverter(DefaultJsonProvider())
         val serialized = sut.serialize(data)
-        assertThat(serialized.size).isGreaterThan(1)
-        assertThat(serialized.size).isEqualTo(100 / WorkerDataConverter.CHUNK_SIZE)
+        assertThat(serialized.getOrNull()?.size).isGreaterThan(1)
+        assertThat(serialized.getOrNull()?.size).isEqualTo(100 / WorkerDataConverter.CHUNK_SIZE)
     }
 
     private fun testIdentity(data: List<NotificationEventRequest>) {
         val sut = WorkerDataConverter(DefaultJsonProvider())
-        val serialized = sut.serialize(data).first().getOrThrow()
-        val result = sut.deserialize(serialized)
+        val serialized = sut.serialize(data).getOrThrow()
+        val result = sut.deserialize(serialized.first())
         assertThat(result).isEqualTo(data)
     }
 }

libraries/workmanager/api/src/main/kotlin/io/element/android/libraries/workmanager/api/WorkManagerRequest.kt

@@ -11,5 +11,5 @@ package io.element.android.libraries.workmanager.api
 import androidx.work.WorkRequest
 
 interface WorkManagerRequest {
-    fun build(): List<Result<WorkRequest>>
+    fun build(): Result<List<WorkRequest>>
 }

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

@@ -27,16 +27,14 @@ class DefaultWorkManagerScheduler(
     private val workManager by lazy { WorkManager.getInstance(context) }
 
     override fun submit(workManagerRequest: WorkManagerRequest) {
-        workManagerRequest.build().forEach {
-            it.fold(
-                onSuccess = { workRequest ->
-                    workManager.enqueue(workRequest)
-                },
-                onFailure = {
-                    Timber.e(it, "Failed to build WorkManager request $workManagerRequest")
-                }
-            )
-        }
+        workManagerRequest.build().fold(
+            onSuccess = { workRequests ->
+                workManager.enqueue(workRequests)
+            },
+            onFailure = {
+                Timber.e(it, "Failed to build WorkManager request $workManagerRequest")
+            }
+        )
     }
 
     override fun cancel(sessionId: SessionId) {