test: TaskManager refactor and tests

spotify · Dec 10, 2024 · 6815649 · 6815649
1 parent 7d4f41f
commit 6815649
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 18 deletions.
diff --git a/Sources/Confidence/Confidence.swift b/Sources/Confidence/Confidence.swift
@@ -23,13 +23,7 @@ public class Confidence: ConfidenceEventSender {
     // Synchronization and task management resources
     private var cancellables = Set<AnyCancellable>()
     private let cacheQueue = DispatchQueue(label: "com.confidence.queue.cache")
-    private var currentFetchTask: Task<(), Never>? {
-        didSet {
-            if let oldTask = oldValue {
-                oldTask.cancel()
-            }
-        }
-    }
+    private var taskManager = TaskManager()
 
     // Internal for testing
     internal let remoteFlagResolver: ConfidenceResolveClient
@@ -161,7 +155,7 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(key: String, value: ConfidenceValue) async {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: [key: value], removedKeys: [])
             do {
                 try await self.fetchAndActivate()
@@ -174,7 +168,7 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(context: ConfidenceStruct, removedKeys: [String] = []) async {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: removedKeys)
             do {
                 try await self.fetchAndActivate()
@@ -187,7 +181,7 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(context: ConfidenceStruct) async {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: [])
             do {
                 try await fetchAndActivate()
@@ -204,7 +198,7 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func removeContextAndWait(key: String) async {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: [:], removedKeys: [key])
             do {
                 try await self.fetchAndActivate()
@@ -231,31 +225,31 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContext(key: String, value: ConfidenceValue) {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             await putContextAndWait(key: key, value: value)
         }
     }
 
     public func putContext(context: ConfidenceStruct) {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             await putContextAndWait(context: context)
         }
     }
 
     public func putContext(context: ConfidenceStruct, removeKeys removedKeys: [String] = []) {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             await putContextAndWait(context: context, removedKeys: removedKeys)
         }
     }
 
     public func removeContext(key: String) {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             await removeContextAndWait(key: key)
         }
     }
 
     public func putContext(context: ConfidenceStruct, removedKeys: [String]) {
-        self.currentFetchTask = Task {
+        taskManager.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: removedKeys)
             do {
                 try await self.fetchAndActivate()
@@ -274,7 +268,7 @@ public class Confidence: ConfidenceEventSender {
     Ensures all the already-started context changes prior to this function have been reconciliated
     */
     public func awaitReconciliation() async {
-        while let task = self.currentFetchTask {
+        while let task = taskManager.currentFetchTask {
             // If current task is cancelled, return
             if task.isCancelled {
                 return
@@ -287,7 +281,7 @@ public class Confidence: ConfidenceEventSender {
             }
             // If current task finished successfully
             // and the set task has not changed, we are done waiting
-            if self.currentFetchTask == task {
+            if taskManager.currentFetchTask == task {
                 return
             }
         }

diff --git a/Sources/Confidence/TaskManager.swift b/Sources/Confidence/TaskManager.swift
@@ -0,0 +1,31 @@
+import Foundation
+
+class TaskManager {
+    public var currentFetchTask: Task<(), Never>? {
+        didSet {
+            if let oldTask = oldValue {
+                print("Cancelling old task")
+                oldTask.cancel()
+            }
+        }
+    }
+    public func awaitReconciliation() async {
+        while let task = self.currentFetchTask {
+            // If current task is cancelled, return
+            if task.isCancelled {
+                return
+            }
+            // Wait for result of current task
+            await task.value
+            // If current task gets cancelled, check again if a new task was set
+            if task.isCancelled {
+                continue
+            }
+            // If current task finished successfully
+            // and the set task has not changed, we are done waiting
+            if self.currentFetchTask == task {
+                return
+            }
+        }
+    }
+}
diff --git a/Tests/ConfidenceTests/TaskManagerTests.swift b/Tests/ConfidenceTests/TaskManagerTests.swift
@@ -0,0 +1,62 @@
+import Foundation
+import XCTest
+@testable import Confidence
+
+class TaskManagerTests: XCTestCase {
+    func testAwaitReconciliationCancelTask() async throws {
+        var signal1 = false
+        let reconciliationExpectation = XCTestExpectation(description: "reconciliationExpectation")
+        let cancelTaskExpectation = XCTestExpectation(description: "cancelTaskExpectation")
+        let taskManager = TaskManager()
+
+        let tenSeconds = Task {
+            do {
+                try await Task.sleep(nanoseconds: 10_000_000_000)
+                signal1 = true
+            } catch {
+                cancelTaskExpectation.fulfill()
+            }
+        }
+        taskManager.currentFetchTask = tenSeconds
+        Task {
+            await taskManager.awaitReconciliation()
+            reconciliationExpectation.fulfill()
+        }
+        tenSeconds.cancel()
+        await fulfillment(of: [cancelTaskExpectation, reconciliationExpectation], timeout: 1)
+        XCTAssertEqual(signal1, false)
+    }
+
+    func testOverrideTask() async throws {
+        var signal1 = false
+        var signal2 = false
+        let reconciliationExpectation = XCTestExpectation(description: "reconciliationExpectation")
+        let cancelTaskExpectation = XCTestExpectation(description: "cancelTaskExpectation")
+        let secondTaskExpectation = XCTestExpectation(description: "secondTaskExpectation")
+        let taskManager = TaskManager()
+
+        let tenSeconds1 = Task {
+            do {
+                try await Task.sleep(nanoseconds: 10_000_000_000)
+                signal1 = true
+            } catch {
+                cancelTaskExpectation.fulfill()
+            }
+        }
+        taskManager.currentFetchTask = tenSeconds1
+
+        let tenSeconds2 = Task {
+            signal2 = true
+            secondTaskExpectation.fulfill()
+        }
+        taskManager.currentFetchTask = tenSeconds2
+
+        Task {
+            await taskManager.awaitReconciliation()
+            reconciliationExpectation.fulfill()
+        }
+        await fulfillment(of: [cancelTaskExpectation, reconciliationExpectation, secondTaskExpectation], timeout: 1)
+        XCTAssertEqual(signal1, false)
+        XCTAssertEqual(signal2, true)
+    }
+}