From f3e39c85f516cbc86311c30d576a07bf3136e9ca Mon Sep 17 00:00:00 2001
From: tyiuhc <tim.yiu@amplitude.com>
Date: Fri, 26 Jan 2024 17:42:23 -0800
Subject: [PATCH] add unit test

---
 Sources/Experiment/ExperimentClient.swift     |  2 +-
 .../ExperimentClientTests.swift               | 48 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/Sources/Experiment/ExperimentClient.swift b/Sources/Experiment/ExperimentClient.swift
index ea47a8d..95c2e8f 100644
--- a/Sources/Experiment/ExperimentClient.swift
+++ b/Sources/Experiment/ExperimentClient.swift
@@ -532,7 +532,7 @@ internal class DefaultExperimentClient : NSObject, ExperimentClient {
         return task
     }
 
-    private func startRetries(user: ExperimentUser, options: FetchOptions?) {
+    internal func startRetries(user: ExperimentUser, options: FetchOptions?) {
         backoffLock.wait()
         defer { backoffLock.signal() }
         self.backoff?.cancel()
diff --git a/Tests/ExperimentTests/ExperimentClientTests.swift b/Tests/ExperimentTests/ExperimentClientTests.swift
index e34bf64..bd8d280 100644
--- a/Tests/ExperimentTests/ExperimentClientTests.swift
+++ b/Tests/ExperimentTests/ExperimentClientTests.swift
@@ -1031,6 +1031,7 @@ class ExperimentClientTests: XCTestCase {
         var mockFetch: (() -> Result<[String: Variant], Error>)? = nil
         var flagCalls = 0
         var mockFlags: (() -> Result<[String: EvaluationFlag], Error>)? = nil
+        var startRetriesCalls = 0
         
         override func doFetch(
             user: ExperimentUser,
@@ -1058,6 +1059,11 @@ class ExperimentClientTests: XCTestCase {
                 super.doFlags(timeoutMillis: timeoutMillis, completion: completion)
             }
         }
+        
+        override func startRetries(user: ExperimentUser, options: FetchOptions?) {
+            startRetriesCalls += 1
+            return
+        }
     }
     
     func testStart_WithLocalAndRemoteEvaluation_CallsFetch() {
@@ -1187,6 +1193,42 @@ class ExperimentClientTests: XCTestCase {
         XCTAssertEqual("on", variant.key!)
         XCTAssertEqual("on", variant2.key!)
     }
+    
+    func testRemoteFetchRetry() {
+        // Response code, error message, and whether retry should be called
+        let testData: [(Int, String, Int)] = [
+            (300, "Fetch Exception 300", 1),
+            (400, "Fetch Exception 400", 0),
+            (429, "Fetch Exception 429", 1),
+            (500, "Fetch Exception 500", 1),
+            (0, "Other Exception", 1)
+        ]
+
+        for (responseCode, errorMessage, retryCalled) in testData {
+            let storage = InMemoryStorage()
+            let config = ExperimentConfigBuilder()
+                .fetchRetryOnFailure(true)
+                .build()
+            let client = MockClient(
+                apiKey: API_KEY,
+                config: config,
+                storage: storage
+            )
+            client.mockFetch = {
+                let error = FetchError(responseCode, errorMessage)
+                return .failure(error)
+            }
+            client.mockFlags = {
+                return .success([:])
+            }
+            let user = ExperimentUserBuilder()
+                .userId("test_retry_fetch")
+                .build()
+            
+            client.fetchBlocking(user: user, isTestRetry: true)
+            XCTAssertEqual(retryCalled, client.startRetriesCalls)
+        }
+    }
 }
 
 class TestAnalyticsProvider : ExperimentAnalyticsProvider {
@@ -1271,11 +1313,13 @@ extension DefaultExperimentClient {
             }
         }
     }
-    func fetchBlocking(user: ExperimentUser) {
+    func fetchBlocking(user: ExperimentUser, isTestRetry: Bool = false) {
         let s = DispatchSemaphore(value: 0)
         fetch(user: user) { _, error in
             if let error = error {
-                XCTFail(error.localizedDescription)
+                if (!isTestRetry) {
+                    XCTFail(error.localizedDescription)
+                }
             }
             s.signal()
         }