fix: Improve remote evaluation fetch retry logic

build.gradle.kts

@@ -19,7 +19,7 @@ java {
 dependencies {
     implementation(kotlin("stdlib"))
     testImplementation(kotlin("test"))
-    testImplementation("org.mockito:mockito-core:${Versions.mockito}")
+    testImplementation("io.mockk:mockk:${Versions.mockk}")
     implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:${Versions.serializationRuntime}")
     implementation("com.squareup.okhttp3:okhttp:${Versions.okhttp}")
     implementation("com.amplitude:evaluation-core:${Versions.evaluationCore}")

buildSrc/src/main/kotlin/Versions.kt

@@ -9,5 +9,5 @@ object Versions {
     const val evaluationCore = "1.1.1"
     const val evaluationSerialization = "1.1.1"
     const val amplitudeAnalytics = "1.12.0"
-    const val mockito = "4.8.1"
+    const val mockk = "1.13.9"
 }

src/main/kotlin/RemoteEvaluationClient.kt

@@ -2,6 +2,7 @@ package com.amplitude.experiment
 
 import com.amplitude.experiment.evaluation.serialization.SerialVariant
 import com.amplitude.experiment.util.BackoffConfig
+import com.amplitude.experiment.util.FetchException
 import com.amplitude.experiment.util.Logger
 import com.amplitude.experiment.util.backoff
 import com.amplitude.experiment.util.toSerialExperimentUser
@@ -44,7 +45,7 @@ class RemoteEvaluationClient internal constructor(
     fun fetch(user: ExperimentUser): CompletableFuture<Map<String, Variant>> {
         return doFetch(user, config.fetchTimeoutMillis).handle { variants, t ->
             if (t != null || variants == null) {
-                if (retry) {
+                if (retry && shouldRetryFetch(t)) {
                     backoff(backoffConfig) {
                         doFetch(user, config.fetchTimeoutMillis)
                     }
@@ -90,7 +91,7 @@ class RemoteEvaluationClient internal constructor(
                     Logger.d("Received fetch response: $response")
                     val variants = response.use {
                         if (!response.isSuccessful) {
-                            throw IOException("fetch error response: $response")
+                            throw FetchException(response.code, "fetch error response: $response")
                         }
                         parseRemoteResponse(response.body?.string() ?: "")
                     }
@@ -112,3 +113,10 @@ internal fun parseRemoteResponse(jsonString: String): Map<String, Variant> =
     json.decodeFromString<HashMap<String, SerialVariant>>(
         jsonString
     ).mapValues { it.value.toVariant() }
+
+private fun shouldRetryFetch(t: Throwable): Boolean {
+    if (t is FetchException) {
+        return t.statusCode < 400 || t.statusCode >= 500 || t.statusCode == 429
+    }
+    return true
+}

src/main/kotlin/util/FetchException.kt

@@ -0,0 +1,8 @@
+package com.amplitude.experiment.util
+
+import okio.IOException
+
+internal class FetchException (
+    val statusCode: Int,
+    message: String
+) : IOException(message)

src/test/kotlin/RemoteEvaluationClientTest.kt

@@ -1,7 +1,11 @@
 package com.amplitude.experiment
 
+import com.amplitude.experiment.util.FetchException
 import com.amplitude.experiment.util.Logger
 import com.amplitude.experiment.util.SystemLogger
+import io.mockk.every
+import io.mockk.spyk
+import io.mockk.verify
 import org.junit.Assert
 import java.util.Date
 import java.util.concurrent.CancellationException
@@ -84,6 +88,7 @@ class RemoteEvaluationClientTest {
         }
     }
 
+
     // @Test
     // fun test() {
     //     val client = RemoteEvaluationClient(API_KEY, RemoteEvaluationConfig())
@@ -100,6 +105,41 @@ class RemoteEvaluationClientTest {
     //         }
     //     }
     // }
+
+    @Test
+    fun `fetch retry with different response codes`() {
+        // Response code, error message, and whether retry should be called
+        val testData = listOf(
+            Triple(300, "Fetch Exception 300", 2),
+            Triple(400, "Fetch Exception 400", 1),
+            Triple(429, "Fetch Exception 429", 2),
+            Triple(500, "Fetch Exception 500", 2),
+            Triple(0, "Other Exception", 2)
+        )
+
+        testData.forEach { (responseCode, errorMessage, fetchCalls) ->
+            val config = RemoteEvaluationConfig(fetchRetries = 1, debug = true)
+            val client = spyk(RemoteEvaluationClient("apiKey", config), recordPrivateCalls = true)
+            // Mock the private method to throw FetchException or other exceptions
+            every { client["doFetch"](any<ExperimentUser>(), any<Long>()) } answers {
+                val future = CompletableFuture<Map<String, Variant>>()
+                if (responseCode == 0) {
+                    future.completeExceptionally(Exception(errorMessage))
+                } else {
+                    future.completeExceptionally(FetchException(responseCode, errorMessage))
+                }
+                future
+            }
+
+            try {
+                client.fetch(ExperimentUser("test_user")).get()
+            } catch (t: Throwable) {
+                // catch exception
+            }
+
+            verify(exactly = fetchCalls) { client["doFetch"](any<ExperimentUser>(), any<Long>()) }
+        }
+    }
 }
 
 @Suppress("SameParameterValue")