Merge pull request #5587 from element-hq/feature/bma/robustWellknown

Improve wellknown retrieval API

features/login/impl/src/main/kotlin/io/element/android/features/login/impl/accesscontrol/DefaultAccountProviderAccessControl.kt

@@ -41,7 +41,7 @@ class DefaultAccountProviderAccessControl(
             // Ensure that Element Pro is not required for this account provider
             val wellKnown = wellknownRetriever.getElementWellKnown(
                 baseUrl = accountProviderUrl.ensureProtocol(),
-            )
+            ).dataOrNull()
             if (wellKnown?.enforceElementPro == true) {
                 throw AccountProviderAccessException.NeedElementProException(
                     unauthorisedAccountProviderTitle = title,

features/login/impl/src/main/kotlin/io/element/android/features/login/impl/resolver/HomeserverResolver.kt

@@ -46,7 +46,7 @@ class HomeserverResolver(
                         wellknownRetriever.getWellKnown(url)
                     }
                 }
-                val isValid = wellKnown?.isValid().orFalse()
+                val isValid = wellKnown?.dataOrNull()?.isValid().orFalse()
                 if (isValid) {
                     // Emit the list as soon as possible
                     currentList.add(

features/login/impl/src/main/kotlin/io/element/android/features/login/impl/web/WebClientUrlForAuthenticationRetriever.kt

@@ -30,7 +30,7 @@ class DefaultWebClientUrlForAuthenticationRetriever(
             Timber.w("Temporary account creation flow is only supported on matrix.org")
             throw AccountCreationNotSupported()
         }
-        val wellknown = wellknownRetriever.getElementWellKnown(homeServerUrl)
+        val wellknown = wellknownRetriever.getElementWellKnown(homeServerUrl).dataOrNull()
             ?: throw AccountCreationNotSupported()
         val registrationHelperUrl = wellknown.registrationHelperUrl
         return if (registrationHelperUrl != null) {

features/login/impl/src/test/kotlin/io/element/android/features/login/impl/accesscontrol/DefaultAccountProviderAccessControlTest.kt

@@ -16,6 +16,7 @@ import io.element.android.libraries.matrix.test.AN_ACCOUNT_PROVIDER
 import io.element.android.libraries.matrix.test.AN_ACCOUNT_PROVIDER_2
 import io.element.android.libraries.matrix.test.AN_ACCOUNT_PROVIDER_URL
 import io.element.android.libraries.wellknown.api.ElementWellKnown
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertThrows
 import org.junit.Test
@@ -155,7 +156,13 @@ class DefaultAccountProviderAccessControlTest {
             defaultHomeserverListResult = { allowedAccountProviders },
         ),
         wellknownRetriever = FakeWellknownRetriever(
-            getElementWellKnownResult = { elementWellKnown },
+            getElementWellKnownResult = {
+                if (elementWellKnown == null) {
+                    WellknownRetrieverResult.NotFound
+                } else {
+                    WellknownRetrieverResult.Success(elementWellKnown)
+                }
+            },
         ),
     )
 

features/login/impl/src/test/kotlin/io/element/android/features/login/impl/changeserver/ChangeServerPresenterTest.kt

@@ -23,6 +23,7 @@ import io.element.android.libraries.matrix.test.A_HOMESERVER_URL
 import io.element.android.libraries.matrix.test.auth.FakeMatrixAuthenticationService
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.WellknownRetriever
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import io.element.android.tests.testutils.WarmUpRule
 import io.element.android.tests.testutils.lambda.lambdaRecorder
 import io.element.android.tests.testutils.lambda.value
@@ -114,9 +115,11 @@ class ChangeServerPresenterTest {
 
     @Test
     fun `present - change server element pro required error`() = runTest {
-        val getElementWellKnownResult = lambdaRecorder<String, ElementWellKnown> {
-            anElementWellKnown(
-                enforceElementPro = true,
+        val getElementWellKnownResult = lambdaRecorder<String, WellknownRetrieverResult<ElementWellKnown>> {
+            WellknownRetrieverResult.Success(
+                anElementWellKnown(
+                    enforceElementPro = true,
+                )
             )
         }
         createPresenter(

features/login/impl/src/test/kotlin/io/element/android/features/login/impl/screens/searchaccountprovider/SearchAccountProviderPresenterTest.kt

@@ -18,6 +18,7 @@ import io.element.android.libraries.architecture.AsyncData
 import io.element.android.libraries.matrix.test.A_HOMESERVER_URL
 import io.element.android.libraries.wellknown.api.WellKnown
 import io.element.android.libraries.wellknown.api.WellKnownBaseConfig
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import io.element.android.tests.testutils.WarmUpRule
 import io.element.android.tests.testutils.lambda.lambdaRecorder
 import io.element.android.tests.testutils.lambda.value
@@ -94,12 +95,12 @@ class SearchAccountProviderPresenterTest {
 
     @Test
     fun `present - enter text one result with wellknown`() = runTest {
-        val getWellKnownResult = lambdaRecorder<String, WellKnown> {
+        val getWellKnownResult = lambdaRecorder<String, WellknownRetrieverResult<WellKnown>> {
             when (it) {
-                "https://test.org" -> error("not found")
-                "https://test.com" -> error("not found")
-                "https://test.io" -> aWellKnown()
-                "https://test" -> error("not found")
+                "https://test.org" -> WellknownRetrieverResult.NotFound
+                "https://test.com" -> WellknownRetrieverResult.NotFound
+                "https://test.io" -> WellknownRetrieverResult.Success(aWellKnown())
+                "https://test" -> WellknownRetrieverResult.NotFound
                 else -> error("should not happen")
             }
         }
@@ -138,12 +139,12 @@ class SearchAccountProviderPresenterTest {
 
     @Test
     fun `present - enter text two results with wellknown`() = runTest {
-        val getWellKnownResult = lambdaRecorder<String, WellKnown> {
+        val getWellKnownResult = lambdaRecorder<String, WellknownRetrieverResult<WellKnown>> {
             when (it) {
-                "https://test.org" -> aWellKnown()
-                "https://test.com" -> error("not found")
-                "https://test.io" -> aWellKnown()
-                "https://test" -> error("not found")
+                "https://test.org" -> WellknownRetrieverResult.Success(aWellKnown())
+                "https://test.com" -> WellknownRetrieverResult.NotFound
+                "https://test.io" -> WellknownRetrieverResult.Success(aWellKnown())
+                "https://test" -> WellknownRetrieverResult.NotFound
                 else -> error("should not happen")
             }
         }

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClient.kt

@@ -286,7 +286,7 @@ class RustMatrixClient(
     override suspend fun getUrl(url: String): Result<ByteArray> = withContext(sessionDispatcher) {
         runCatchingExceptions {
             innerClient.getUrl(url)
-        }
+        }.mapFailure { it.mapClientException() }
     }
 
     override suspend fun getRoom(roomId: RoomId): BaseRoom? = withContext(sessionDispatcher) {

libraries/wellknown/api/src/main/kotlin/io/element/android/libraries/wellknown/api/SessionWellknownRetriever.kt

@@ -8,6 +8,6 @@
 package io.element.android.libraries.wellknown.api
 
 interface SessionWellknownRetriever {
-    suspend fun getWellKnown(): WellKnown?
-    suspend fun getElementWellKnown(): ElementWellKnown?
+    suspend fun getWellKnown(): WellknownRetrieverResult<WellKnown>
+    suspend fun getElementWellKnown(): WellknownRetrieverResult<ElementWellKnown>
 }

libraries/wellknown/api/src/main/kotlin/io/element/android/libraries/wellknown/api/WellknownRetriever.kt

@@ -8,6 +8,6 @@
 package io.element.android.libraries.wellknown.api
 
 interface WellknownRetriever {
-    suspend fun getWellKnown(baseUrl: String): WellKnown?
-    suspend fun getElementWellKnown(baseUrl: String): ElementWellKnown?
+    suspend fun getWellKnown(baseUrl: String): WellknownRetrieverResult<WellKnown>
+    suspend fun getElementWellKnown(baseUrl: String): WellknownRetrieverResult<ElementWellKnown>
 }

libraries/wellknown/api/src/main/kotlin/io/element/android/libraries/wellknown/api/WellknownRetrieverResult.kt

@@ -0,0 +1,31 @@
+/*
+ * 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.wellknown.api
+
+sealed interface WellknownRetrieverResult<out T> {
+    /**
+     * Well-known data has been successfully retrieved.
+     */
+    data class Success<out T>(val data: T) : WellknownRetrieverResult<T>
+
+    /**
+     * Well-known data is not found (file does not exist server side, we got a 404).
+     */
+    data object NotFound : WellknownRetrieverResult<Nothing>
+
+    /**
+     * Any other error.
+     */
+    data class Error(val exception: Exception) : WellknownRetrieverResult<Nothing>
+
+    fun dataOrNull(): T? = when (this) {
+        is Success<T> -> data
+        is Error -> null
+        NotFound -> null
+    }
+}

libraries/wellknown/impl/src/main/kotlin/io/element/android/libraries/wellknown/impl/DefaultSessionWellknownRetriever.kt

@@ -13,9 +13,11 @@ import io.element.android.libraries.androidutils.json.JsonProvider
 import io.element.android.libraries.core.extensions.mapCatchingExceptions
 import io.element.android.libraries.di.SessionScope
 import io.element.android.libraries.matrix.api.MatrixClient
+import io.element.android.libraries.matrix.api.exception.ClientException
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.SessionWellknownRetriever
 import io.element.android.libraries.wellknown.api.WellKnown
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import timber.log.Timber
 
 @ContributesBinding(SessionScope::class)
@@ -26,29 +28,40 @@ class DefaultSessionWellknownRetriever(
 ) : SessionWellknownRetriever {
     private val domain by lazy { matrixClient.userIdServerName() }
 
-    override suspend fun getWellKnown(): WellKnown? {
+    override suspend fun getWellKnown(): WellknownRetrieverResult<WellKnown> {
         val url = "https://$domain/.well-known/matrix/client"
         return matrixClient
             .getUrl(url)
             .mapCatchingExceptions {
                 val data = String(it)
-                json().decodeFromString(InternalWellKnown.serializer(), data)
+                json().decodeFromString<InternalWellKnown>(data).map()
             }
-            .onFailure { Timber.e(it, "Failed to retrieve .well-known from $domain") }
-            .map { it.map() }
-            .getOrNull()
+            .toWellknownRetrieverResult()
     }
 
-    override suspend fun getElementWellKnown(): ElementWellKnown? {
+    override suspend fun getElementWellKnown(): WellknownRetrieverResult<ElementWellKnown> {
         val url = "https://$domain/.well-known/element/element.json"
         return matrixClient
             .getUrl(url)
             .mapCatchingExceptions {
                 val data = String(it)
-                json().decodeFromString(InternalElementWellKnown.serializer(), data)
+                json().decodeFromString<InternalElementWellKnown>(data).map()
             }
-            .onFailure { Timber.e(it, "Failed to retrieve Element .well-known from $domain") }
-            .map { it.map() }
-            .getOrNull()
+            .toWellknownRetrieverResult()
     }
+
+    private fun <T> Result<T>.toWellknownRetrieverResult(): WellknownRetrieverResult<T> = fold(
+        onSuccess = {
+            WellknownRetrieverResult.Success(it)
+        },
+        onFailure = {
+            Timber.e(it, "Failed to retrieve Element .well-known from $domain")
+            // This check on message value is not ideal but this is what we got from the SDK.
+            if ((it as? ClientException.Generic)?.message?.contains("404") == true) {
+                WellknownRetrieverResult.NotFound
+            } else {
+                WellknownRetrieverResult.Error(it as Exception)
+            }
+        }
+    )
 }

libraries/wellknown/impl/src/main/kotlin/io/element/android/libraries/wellknown/impl/DefaultWellknownRetriever.kt

@@ -10,46 +10,72 @@ package io.element.android.libraries.wellknown.impl
 import dev.zacsweers.metro.AppScope
 import dev.zacsweers.metro.ContributesBinding
 import dev.zacsweers.metro.Inject
+import io.element.android.libraries.core.extensions.runCatchingExceptions
 import io.element.android.libraries.core.uri.ensureProtocol
 import io.element.android.libraries.network.RetrofitFactory
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.WellKnown
 import io.element.android.libraries.wellknown.api.WellknownRetriever
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
+import retrofit2.HttpException
 import timber.log.Timber
+import java.net.HttpURLConnection
 
 @ContributesBinding(AppScope::class)
 @Inject
 class DefaultWellknownRetriever(
     private val retrofitFactory: RetrofitFactory,
 ) : WellknownRetriever {
-    override suspend fun getWellKnown(baseUrl: String): WellKnown? {
-        val wellknownApi = buildWellknownApi(baseUrl) ?: return null
-        return try {
-            wellknownApi.getWellKnown().map()
-        } catch (e: Exception) {
-            Timber.e(e, "Failed to retrieve well-known data for $baseUrl")
-            null
-        }
+    override suspend fun getWellKnown(baseUrl: String): WellknownRetrieverResult<WellKnown> {
+        return buildWellknownApi(baseUrl)
+            .map { wellknownApi ->
+                try {
+                    val result = wellknownApi.getWellKnown().map()
+                    WellknownRetrieverResult.Success(result)
+                } catch (e: Exception) {
+                    Timber.e(e, "Failed to retrieve well-known data for $baseUrl")
+                    if ((e as? HttpException)?.code() == HttpURLConnection.HTTP_NOT_FOUND) {
+                        WellknownRetrieverResult.NotFound
+                    } else {
+                        WellknownRetrieverResult.Error(e)
+                    }
+                }
+            }
+            .fold(
+                onSuccess = { it },
+                onFailure = { WellknownRetrieverResult.Error(it as Exception) }
+            )
     }
 
-    override suspend fun getElementWellKnown(baseUrl: String): ElementWellKnown? {
-        val wellknownApi = buildWellknownApi(baseUrl) ?: return null
-        return try {
-            wellknownApi.getElementWellKnown().map()
-        } catch (e: Exception) {
-            Timber.e(e, "Failed to retrieve Element well-known data for $baseUrl")
-            null
-        }
+    override suspend fun getElementWellKnown(baseUrl: String): WellknownRetrieverResult<ElementWellKnown> {
+        return buildWellknownApi(baseUrl)
+            .map { wellknownApi ->
+                try {
+                    val result = wellknownApi.getElementWellKnown().map()
+                    WellknownRetrieverResult.Success(result)
+                } catch (e: Exception) {
+                    // Is it a 404?
+                    Timber.e(e, "Failed to retrieve Element well-known data for $baseUrl")
+                    if ((e as? HttpException)?.code() == HttpURLConnection.HTTP_NOT_FOUND) {
+                        WellknownRetrieverResult.NotFound
+                    } else {
+                        WellknownRetrieverResult.Error(e)
+                    }
+                }
+            }
+            .fold(
+                onSuccess = { it },
+                onFailure = { WellknownRetrieverResult.Error(it as Exception) }
+            )
     }
 
-    private fun buildWellknownApi(accountProviderUrl: String): WellknownAPI? {
-        return try {
+    private fun buildWellknownApi(accountProviderUrl: String): Result<WellknownAPI> {
+        return runCatchingExceptions {
             retrofitFactory.create(accountProviderUrl.ensureProtocol())
                 .create(WellknownAPI::class.java)
-        } catch (e: Exception) {
+        }.onFailure { e ->
             // If the base URL is not valid, we cannot retrieve the well-known data
             Timber.e(e, "Failed to create Retrofit instance for $accountProviderUrl")
-            null
         }
     }
 }

libraries/wellknown/impl/src/test/kotlin/io/element/android/libraries/wellknown/impl/DefaultSessionWellknownRetrieverTest.kt

@@ -14,6 +14,7 @@ import io.element.android.libraries.matrix.test.FakeMatrixClient
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.WellKnown
 import io.element.android.libraries.wellknown.api.WellKnownBaseConfig
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import io.element.android.tests.testutils.lambda.lambdaRecorder
 import io.element.android.tests.testutils.lambda.value
 import kotlinx.coroutines.test.runTest
@@ -29,9 +30,11 @@ class DefaultSessionWellknownRetrieverTest {
             getUrlLambda = getUrlLambda,
         )
         assertThat(sut.getWellKnown()).isEqualTo(
-            WellKnown(
-                homeServer = null,
-                identityServer = null,
+            WellknownRetrieverResult.Success(
+                WellKnown(
+                    homeServer = null,
+                    identityServer = null,
+                )
             )
         )
         getUrlLambda.assertions().isCalledOnce()
@@ -55,13 +58,15 @@ class DefaultSessionWellknownRetrieverTest {
             }
         )
         assertThat(sut.getWellKnown()).isEqualTo(
-            WellKnown(
-                homeServer = WellKnownBaseConfig(
-                    baseURL = "https://example.org",
-                ),
-                identityServer = WellKnownBaseConfig(
-                    baseURL = "https://identity.example.org",
-                ),
+            WellknownRetrieverResult.Success(
+                WellKnown(
+                    homeServer = WellKnownBaseConfig(
+                        baseURL = "https://example.org",
+                    ),
+                    identityServer = WellKnownBaseConfig(
+                        baseURL = "https://identity.example.org",
+                    ),
+                )
             )
         )
     }
@@ -81,13 +86,15 @@ class DefaultSessionWellknownRetrieverTest {
             }
         )
         assertThat(sut.getWellKnown()).isEqualTo(
-            WellKnown(
-                homeServer = WellKnownBaseConfig(
-                    baseURL = "https://example.org",
-                ),
-                identityServer = WellKnownBaseConfig(
-                    baseURL = null,
-                ),
+            WellknownRetrieverResult.Success(
+                WellKnown(
+                    homeServer = WellKnownBaseConfig(
+                        baseURL = "https://example.org",
+                    ),
+                    identityServer = WellKnownBaseConfig(
+                        baseURL = null,
+                    ),
+                )
             )
         )
     }
@@ -110,13 +117,15 @@ class DefaultSessionWellknownRetrieverTest {
             },
         )
         assertThat(sut.getWellKnown()).isEqualTo(
-            WellKnown(
-                homeServer = WellKnownBaseConfig(
-                    baseURL = "https://example.org",
-                ),
-                identityServer = WellKnownBaseConfig(
-                    baseURL = "https://identity.example.org",
-                ),
+            WellknownRetrieverResult.Success(
+                WellKnown(
+                    homeServer = WellKnownBaseConfig(
+                        baseURL = "https://example.org",
+                    ),
+                    identityServer = WellKnownBaseConfig(
+                        baseURL = "https://identity.example.org",
+                    ),
+                )
             )
         )
     }
@@ -135,7 +144,7 @@ class DefaultSessionWellknownRetrieverTest {
                 )
             }
         )
-        assertThat(sut.getWellKnown()).isNull()
+        assertThat(sut.getWellKnown()).isInstanceOf(WellknownRetrieverResult.Error::class.java)
     }
 
     @Test
@@ -145,7 +154,7 @@ class DefaultSessionWellknownRetrieverTest {
                 Result.failure(AN_EXCEPTION)
             }
         )
-        assertThat(sut.getWellKnown()).isNull()
+        assertThat(sut.getWellKnown()).isInstanceOf(WellknownRetrieverResult.Error::class.java)
     }
 
     @Test
@@ -157,11 +166,13 @@ class DefaultSessionWellknownRetrieverTest {
             getUrlLambda = getUrlLambda,
         )
         assertThat(sut.getElementWellKnown()).isEqualTo(
-            ElementWellKnown(
-                registrationHelperUrl = null,
-                enforceElementPro = null,
-                rageshakeUrl = null,
-                brandColor = null,
+            WellknownRetrieverResult.Success(
+                ElementWellKnown(
+                    registrationHelperUrl = null,
+                    enforceElementPro = null,
+                    rageshakeUrl = null,
+                    brandColor = null,
+                )
             )
         )
         getUrlLambda.assertions().isCalledOnce()
@@ -183,11 +194,13 @@ class DefaultSessionWellknownRetrieverTest {
             }
         )
         assertThat(sut.getElementWellKnown()).isEqualTo(
-            ElementWellKnown(
-                registrationHelperUrl = "a_registration_url",
-                enforceElementPro = true,
-                rageshakeUrl = "a_rageshake_url",
-                brandColor = "#FF0000",
+            WellknownRetrieverResult.Success(
+                ElementWellKnown(
+                    registrationHelperUrl = "a_registration_url",
+                    enforceElementPro = true,
+                    rageshakeUrl = "a_rageshake_url",
+                    brandColor = "#FF0000",
+                )
             )
         )
     }
@@ -207,11 +220,13 @@ class DefaultSessionWellknownRetrieverTest {
             },
         )
         assertThat(sut.getElementWellKnown()).isEqualTo(
-            ElementWellKnown(
-                registrationHelperUrl = "a_registration_url",
-                enforceElementPro = true,
-                rageshakeUrl = "a_rageshake_url",
-                brandColor = null,
+            WellknownRetrieverResult.Success(
+                ElementWellKnown(
+                    registrationHelperUrl = "a_registration_url",
+                    enforceElementPro = true,
+                    rageshakeUrl = "a_rageshake_url",
+                    brandColor = null,
+                )
             )
         )
     }
@@ -228,7 +243,7 @@ class DefaultSessionWellknownRetrieverTest {
                 )
             }
         )
-        assertThat(sut.getElementWellKnown()).isNull()
+        assertThat(sut.getElementWellKnown()).isInstanceOf(WellknownRetrieverResult.Error::class.java)
     }
 
     @Test
@@ -238,7 +253,7 @@ class DefaultSessionWellknownRetrieverTest {
                 Result.failure(AN_EXCEPTION)
             }
         )
-        assertThat(sut.getElementWellKnown()).isNull()
+        assertThat(sut.getElementWellKnown()).isInstanceOf(WellknownRetrieverResult.Error::class.java)
     }
 
     private fun createDefaultSessionWellknownRetriever(

libraries/wellknown/test/src/main/kotlin/io/element/android/features/wellknown/test/FakeSessionWellknownRetriever.kt

@@ -10,17 +10,18 @@ package io.element.android.features.wellknown.test
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.SessionWellknownRetriever
 import io.element.android.libraries.wellknown.api.WellKnown
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import io.element.android.tests.testutils.simulateLongTask
 
 class FakeSessionWellknownRetriever(
-    private val getWellKnownResult: () -> WellKnown? = { null },
-    private val getElementWellKnownResult: () -> ElementWellKnown? = { null },
+    private val getWellKnownResult: () -> WellknownRetrieverResult<WellKnown> = { WellknownRetrieverResult.NotFound },
+    private val getElementWellKnownResult: () -> WellknownRetrieverResult<ElementWellKnown> = { WellknownRetrieverResult.NotFound },
 ) : SessionWellknownRetriever {
-    override suspend fun getWellKnown(): WellKnown? = simulateLongTask {
+    override suspend fun getWellKnown(): WellknownRetrieverResult<WellKnown> = simulateLongTask {
         getWellKnownResult()
     }
 
-    override suspend fun getElementWellKnown(): ElementWellKnown? = simulateLongTask {
+    override suspend fun getElementWellKnown(): WellknownRetrieverResult<ElementWellKnown> = simulateLongTask {
         getElementWellKnownResult()
     }
 }

libraries/wellknown/test/src/main/kotlin/io/element/android/features/wellknown/test/FakeWellknownRetriever.kt

@@ -10,17 +10,18 @@ package io.element.android.features.wellknown.test
 import io.element.android.libraries.wellknown.api.ElementWellKnown
 import io.element.android.libraries.wellknown.api.WellKnown
 import io.element.android.libraries.wellknown.api.WellknownRetriever
+import io.element.android.libraries.wellknown.api.WellknownRetrieverResult
 import io.element.android.tests.testutils.simulateLongTask
 
 class FakeWellknownRetriever(
-    private val getWellKnownResult: (String) -> WellKnown? = { null },
-    private val getElementWellKnownResult: (String) -> ElementWellKnown? = { null },
+    private val getWellKnownResult: (String) -> WellknownRetrieverResult<WellKnown> = { WellknownRetrieverResult.NotFound },
+    private val getElementWellKnownResult: (String) -> WellknownRetrieverResult<ElementWellKnown> = { WellknownRetrieverResult.NotFound },
 ) : WellknownRetriever {
-    override suspend fun getWellKnown(baseUrl: String): WellKnown? = simulateLongTask {
+    override suspend fun getWellKnown(baseUrl: String): WellknownRetrieverResult<WellKnown> = simulateLongTask {
         getWellKnownResult(baseUrl)
     }
 
-    override suspend fun getElementWellKnown(baseUrl: String): ElementWellKnown? = simulateLongTask {
+    override suspend fun getElementWellKnown(baseUrl: String): WellknownRetrieverResult<ElementWellKnown> = simulateLongTask {
         getElementWellKnownResult(baseUrl)
     }
 }