Read install referrer on Dispatchers.IO

ET-6180
2026-05-15 09:50:40 +00:00 · 2026-05-06 09:54:05 +02:00
parent 0d61db6090
commit 8a142815a3
4 changed files with 243 additions and 51 deletions
@@ -51,7 +51,9 @@ dependencies {
    kapt(libs.bundles.app.annotationProcessors)
    implementation(libs.dagger.hilt.android)

+    implementation(project(":mail-common:domain"))
    implementation(project(":mail-events:data"))
    implementation(project(":mail-events:domain"))
    implementation(project(":shared:core:events:domain"))
+    implementation(libs.kotlin.coroutines.core)
 }
@@ -54,11 +54,6 @@ object MailEventsModule {
    fun provideEventsDataStoreProvider(@ApplicationContext context: Context): EventsDataStoreProvider =
        EventsDataStoreProvider(context)

-    @Provides
-    @Singleton
-    fun provideInstallReferrerDataSource(@ApplicationContext context: Context): InstallReferrerDataSource =
-        PlayInstallReferrerDataSourceImpl(context)
-
    @Module
    @InstallIn(SingletonComponent::class)
    internal interface BindsModule {
@@ -79,10 +74,13 @@ object MailEventsModule {
        @Reusable
        fun bindsDeviceInfoProvider(impl: DeviceInfoProviderImpl): DeviceInfoProvider

-
        @Binds
        @Reusable
        fun bindsAppInstallRepository(impl: AppInstallRepositoryImpl): AppInstallRepository
+
+        @Binds
+        @Singleton
+        fun bindsInstallReferrerDataSource(impl: PlayInstallReferrerDataSourceImpl): InstallReferrerDataSource
    }
 }

@@ -22,72 +22,86 @@ import android.content.Context
 import arrow.core.Either
 import arrow.core.left
 import arrow.core.right
+import ch.protonmail.android.mailcommon.domain.coroutines.IODispatcher
 import ch.protonmail.android.mailcommon.domain.model.DataError
 import ch.protonmail.android.mailevents.domain.model.InstallReferrer
 import com.android.installreferrer.api.InstallReferrerClient
 import com.android.installreferrer.api.InstallReferrerStateListener
+import dagger.hilt.android.qualifiers.ApplicationContext
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.withTimeoutOrNull
 import timber.log.Timber
 import javax.inject.Inject
 import kotlin.coroutines.resume
+import kotlin.time.Duration.Companion.milliseconds

@Suppress("TooGenericExceptionCaught")
 class PlayInstallReferrerDataSourceImpl @Inject constructor(
-    private val context: Context
+    @ApplicationContext private val context: Context,
+    @IODispatcher private val ioDispatcher: CoroutineDispatcher
 ) : InstallReferrerDataSource {

-    override suspend fun getInstallReferrer(): Either<DataError, InstallReferrer> =
-        suspendCancellableCoroutine { continuation ->
-            val referrerClient = InstallReferrerClient.newBuilder(context).build()
+    override suspend fun getInstallReferrer(): Either<DataError, InstallReferrer> = withContext(ioDispatcher) {
+        val referrerClient = InstallReferrerClient.newBuilder(context).build()
+        try {
+            val responseCode = withTimeoutOrNull(CONNECTION_TIMEOUT_MS.milliseconds) {
+                awaitConnection(referrerClient)
+            } ?: run {
+                Timber.d("Install referrer connection timed out")
+                return@withContext DataError.Local.Unknown.left()
+            }

+            when (responseCode) {
+                InstallReferrerClient.InstallReferrerResponse.OK -> {
+                    val response = referrerClient.installReferrer
+                    InstallReferrer(
+                        referrerUrl = response.installReferrer,
+                        referrerClickTimestampMs = response.referrerClickTimestampSeconds * 1000,
+                        installBeginTimestampMs = response.installBeginTimestampSeconds * 1000,
+                        isGooglePlayInstant = response.googlePlayInstantParam
+                    ).right()
+                }
+
+                InstallReferrerClient.InstallReferrerResponse.FEATURE_NOT_SUPPORTED -> {
+                    Timber.d("Install referrer not supported on this device")
+                    DataError.Local.Unknown.left()
+                }
+
+                InstallReferrerClient.InstallReferrerResponse.SERVICE_UNAVAILABLE -> {
+                    Timber.d("Install referrer service unavailable")
+                    DataError.Local.Unknown.left()
+                }
+
+                else -> {
+                    Timber.d("Unknown install referrer response code: $responseCode")
+                    DataError.Local.Unknown.left()
+                }
+            }
+        } catch (e: Exception) {
+            Timber.e(e, "Failed to get install referrer details")
+            DataError.Local.Unknown.left()
+        } finally {
+            referrerClient.endConnection()
+        }
+    }
+
+    private suspend fun awaitConnection(referrerClient: InstallReferrerClient): Int =
+        suspendCancellableCoroutine { continuation ->
            referrerClient.startConnection(object : InstallReferrerStateListener {
                override fun onInstallReferrerSetupFinished(responseCode: Int) {
-                    when (responseCode) {
-                        InstallReferrerClient.InstallReferrerResponse.OK -> {
-                            try {
-                                val response = referrerClient.installReferrer
-                                val referrer = InstallReferrer(
-                                    referrerUrl = response.installReferrer,
-                                    referrerClickTimestampMs = response.referrerClickTimestampSeconds * 1000,
-                                    installBeginTimestampMs = response.installBeginTimestampSeconds * 1000,
-                                    isGooglePlayInstant = response.googlePlayInstantParam
-                                )
-                                continuation.resume(referrer.right())
-                            } catch (e: Exception) {
-                                Timber.e(e, "Failed to get install referrer details")
-                                continuation.resume(DataError.Local.Unknown.left())
-                            } finally {
-                                referrerClient.endConnection()
-                            }
-                        }
-
-                        InstallReferrerClient.InstallReferrerResponse.FEATURE_NOT_SUPPORTED -> {
-                            Timber.d("Install referrer not supported on this device")
-                            continuation.resume(DataError.Local.Unknown.left())
-                            referrerClient.endConnection()
-                        }
-
-                        InstallReferrerClient.InstallReferrerResponse.SERVICE_UNAVAILABLE -> {
-                            Timber.d("Install referrer service unavailable")
-                            continuation.resume(DataError.Local.Unknown.left())
-                            referrerClient.endConnection()
-                        }
-
-                        else -> {
-                            Timber.d("Unknown install referrer response code: $responseCode")
-                            continuation.resume(DataError.Local.Unknown.left())
-                            referrerClient.endConnection()
-                        }
-                    }
+                    if (continuation.isActive) continuation.resume(responseCode)
                }

                override fun onInstallReferrerServiceDisconnected() {
                    Timber.d("Install referrer service disconnected")
                }
            })
-
-            continuation.invokeOnCancellation {
-                referrerClient.endConnection()
-            }
        }
+
+    private companion object {
+
+        const val CONNECTION_TIMEOUT_MS = 10_000L
+    }
 }
@@ -0,0 +1,178 @@
+/*
+ * Copyright (c) 2026 Proton Technologies AG
+ * This file is part of Proton Technologies AG and Proton Mail.
+ *
+ * Proton Mail is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Proton Mail is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Proton Mail. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package ch.protonmail.android.mailevents.data.referrer
+
+import android.content.Context
+import arrow.core.left
+import arrow.core.right
+import ch.protonmail.android.mailcommon.domain.model.DataError
+import ch.protonmail.android.mailevents.domain.model.InstallReferrer
+import com.android.installreferrer.api.InstallReferrerClient
+import com.android.installreferrer.api.InstallReferrerStateListener
+import com.android.installreferrer.api.ReferrerDetails
+import io.mockk.Runs
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.slot
+import io.mockk.unmockkStatic
+import io.mockk.verify
+import kotlinx.coroutines.async
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.runTest
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.time.Duration.Companion.milliseconds
+
+internal class PlayInstallReferrerDataSourceImplTest {
+
+    private val context = mockk<Context>()
+    private val builder = mockk<InstallReferrerClient.Builder>()
+    private val client = mockk<InstallReferrerClient>(relaxUnitFun = true)
+    private val listenerSlot = slot<InstallReferrerStateListener>()
+
+    private val testDispatcher = StandardTestDispatcher()
+    private val dataSource = PlayInstallReferrerDataSourceImpl(context, testDispatcher)
+
+    @BeforeTest
+    fun setUp() {
+        mockkStatic(InstallReferrerClient::class)
+        every { InstallReferrerClient.newBuilder(context) } returns builder
+        every { builder.build() } returns client
+        every { client.startConnection(capture(listenerSlot)) } just Runs
+        every { client.endConnection() } just Runs
+    }
+
+    @AfterTest
+    fun tearDown() {
+        unmockkStatic(InstallReferrerClient::class)
+    }
+
+    @Test
+    fun `should return install referrer when response is OK`() = runTest(testDispatcher) {
+        // Given
+        val details = mockk<ReferrerDetails> {
+            every { installReferrer } returns "utm_source=test"
+            every { referrerClickTimestampSeconds } returns 1_000L
+            every { installBeginTimestampSeconds } returns 2_000L
+            every { googlePlayInstantParam } returns false
+        }
+        every { client.installReferrer } returns details
+        every { client.startConnection(capture(listenerSlot)) } answers {
+            listenerSlot.captured.onInstallReferrerSetupFinished(InstallReferrerClient.InstallReferrerResponse.OK)
+        }
+
+        // When
+        val result = dataSource.getInstallReferrer()
+
+        // Then
+        val expected = InstallReferrer(
+            referrerUrl = "utm_source=test",
+            referrerClickTimestampMs = 1_000_000L,
+            installBeginTimestampMs = 2_000_000L,
+            isGooglePlayInstant = false
+        )
+        assertEquals(expected.right(), result)
+        verify { client.endConnection() }
+    }
+
+    @Test
+    fun `should return error when feature is not supported`() = runTest(testDispatcher) {
+        // Given
+        every { client.startConnection(capture(listenerSlot)) } answers {
+            listenerSlot.captured.onInstallReferrerSetupFinished(
+                InstallReferrerClient.InstallReferrerResponse.FEATURE_NOT_SUPPORTED
+            )
+        }
+
+        // When
+        val result = dataSource.getInstallReferrer()
+
+        // Then
+        assertEquals(DataError.Local.Unknown.left(), result)
+        verify { client.endConnection() }
+    }
+
+    @Test
+    fun `should return error when service is unavailable`() = runTest(testDispatcher) {
+        // Given
+        every { client.startConnection(capture(listenerSlot)) } answers {
+            listenerSlot.captured.onInstallReferrerSetupFinished(
+                InstallReferrerClient.InstallReferrerResponse.SERVICE_UNAVAILABLE
+            )
+        }
+
+        // When
+        val result = dataSource.getInstallReferrer()
+
+        // Then
+        assertEquals(DataError.Local.Unknown.left(), result)
+        verify { client.endConnection() }
+    }
+
+    @Test
+    fun `should return error for unknown response code`() = runTest(testDispatcher) {
+        // Given
+        every { client.startConnection(capture(listenerSlot)) } answers {
+            listenerSlot.captured.onInstallReferrerSetupFinished(42)
+        }
+
+        // When
+        val result = dataSource.getInstallReferrer()
+
+        // Then
+        assertEquals(DataError.Local.Unknown.left(), result)
+        verify { client.endConnection() }
+    }
+
+    @Test
+    fun `should return error when reading install referrer throws`() = runTest(testDispatcher) {
+        // Given
+        every { client.installReferrer } throws RuntimeException("exception")
+        every { client.startConnection(capture(listenerSlot)) } answers {
+            listenerSlot.captured.onInstallReferrerSetupFinished(InstallReferrerClient.InstallReferrerResponse.OK)
+        }
+
+        // When
+        val result = dataSource.getInstallReferrer()
+
+        // Then
+        assertEquals(DataError.Local.Unknown.left(), result)
+        verify { client.endConnection() }
+    }
+
+    @Test
+    fun `should return error when connection times out`() = runTest(testDispatcher) {
+        // Given - startConnection never invokes the listener
+        every { client.startConnection(any()) } just Runs
+
+        // When
+        val deferred = async { dataSource.getInstallReferrer() }
+        advanceTimeBy(11_000L.milliseconds)
+        val result = deferred.await()
+
+        // Then
+        assertEquals(DataError.Local.Unknown.left(), result)
+        verify { client.endConnection() }
+    }
+}