Fix flaky tests

kean · May 4, 2024 · 7dc8de8 · 7dc8de8
1 parent 09f3620
commit 7dc8de8
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 32 deletions.
diff --git a/Sources/Nuke/Pipeline/ImagePipeline.swift b/Sources/Nuke/Pipeline/ImagePipeline.swift
@@ -53,6 +53,7 @@ public final class ImagePipeline: @unchecked Sendable {
 
     let rateLimiter: RateLimiter?
     let id = UUID()
+    var onTaskStarted: ((ImageTask) -> Void)? // Debug purposes
 
     deinit {
         lock.deinitialize(count: 1)
@@ -309,6 +310,7 @@ public final class ImagePipeline: @unchecked Sendable {
             task?.process(.init($0))
         }
         delegate.imageTaskDidStart(task, pipeline: self)
+        onTaskStarted?(task)
     }
 
     // MARK: - Image Task Events

diff --git a/Tests/NukeTests/ImagePipelineTests/ImagePipelineCoalescingTests.swift b/Tests/NukeTests/ImagePipelineTests/ImagePipelineCoalescingTests.swift
@@ -488,18 +488,11 @@ class ImagePipelineCoalescingTests: XCTestCase {
             $0.isTaskCoalescingEnabled = false
         }
 
-        dataLoader.queue.isSuspended = true
-
         // When/Then
-        let request1 = Test.request
-        let request2 = Test.request
-
-        expect(pipeline).toLoadImage(with: request1)
-        expect(pipeline).toLoadImage(with: request2)
-
-        pipeline.queue.sync {}
-        dataLoader.queue.isSuspended = false
-
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
+            expect(pipeline).toLoadImage(with: Test.request)
+            expect(pipeline).toLoadImage(with: Test.request)
+        }
         wait { _ in
             XCTAssertEqual(self.dataLoader.createdTaskCount, 2)
         }
@@ -658,16 +651,15 @@ class ImagePipelineProcessingDeduplicationTests: XCTestCase {
         pipeline = pipeline.reconfigured {
             $0.dataCache = dataCache
         }
-        dataLoader.queue.isSuspended = true
 
         // When
         func makeRequest(options: ImageRequest.Options) -> ImageRequest {
             ImageRequest(urlRequest: URLRequest(url: Test.url), options: options)
         }
-        expect(pipeline).toLoadImage(with: makeRequest(options: []))
-        expect(pipeline).toLoadImage(with: makeRequest(options: [.reloadIgnoringCachedData]))
-        pipeline.queue.sync {}
-        dataLoader.queue.isSuspended = false
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
+            expect(pipeline).toLoadImage(with: makeRequest(options: []))
+            expect(pipeline).toLoadImage(with: makeRequest(options: [.reloadIgnoringCachedData]))
+        }
 
         // Then
         wait { _ in
@@ -681,17 +673,17 @@ class ImagePipelineProcessingDeduplicationTests: XCTestCase {
         pipeline = pipeline.reconfigured {
             $0.dataCache = dataCache
         }
-        dataLoader.queue.isSuspended = true
 
         // When
         // - One request reloading cache data, another one not
         func makeRequest(options: ImageRequest.Options) -> ImageRequest {
             ImageRequest(urlRequest: URLRequest(url: Test.url), options: options)
         }
-        expect(pipeline).toLoadImage(with: makeRequest(options: []))
-        expect(pipeline).toLoadImage(with: makeRequest(options: [.reloadIgnoringCachedData]))
-        pipeline.queue.sync {}
-        dataLoader.queue.isSuspended = false
+
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
+            expect(pipeline).toLoadImage(with: makeRequest(options: []))
+            expect(pipeline).toLoadImage(with: makeRequest(options: [.reloadIgnoringCachedData]))
+        }
 
         // Then
         wait { _ in

diff --git a/Tests/NukeTests/ImagePipelineTests/ImagePipelineDataCacheTests.swift b/Tests/NukeTests/ImagePipelineTests/ImagePipelineDataCacheTests.swift
@@ -392,7 +392,7 @@ class ImagePipelineDataCachePolicyTests: XCTestCase {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url))
         }
@@ -479,7 +479,7 @@ class ImagePipelineDataCachePolicyTests: XCTestCase {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url))
         }
@@ -544,7 +544,7 @@ class ImagePipelineDataCachePolicyTests: XCTestCase {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadImage(with: ImageRequest(url: Test.url))
         }

diff --git a/Tests/NukeTests/ImagePipelineTests/ImagePipelineLoadDataTests.swift b/Tests/NukeTests/ImagePipelineTests/ImagePipelineLoadDataTests.swift
@@ -211,7 +211,7 @@ extension ImagePipelineLoadDataTests {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url))
         }
@@ -275,7 +275,7 @@ extension ImagePipelineLoadDataTests {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url))
         }
@@ -401,7 +401,7 @@ extension ImagePipelineLoadDataTests {
         }
 
         // WHEN
-        pipeline.registerMultipleRequests {
+        suspendDataLoading(for: pipeline, expectedRequestCount: 2) {
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url, processors: [MockImageProcessor(id: "p1")]))
             expect(pipeline).toLoadData(with: ImageRequest(url: Test.url))
         }
@@ -417,11 +417,17 @@ extension ImagePipelineLoadDataTests {
     }
 }
 
-extension ImagePipeline {
-    func registerMultipleRequests(_ closure: () -> Void) {
-        configuration.dataLoadingQueue.isSuspended = true
+extension XCTestCase {
+    func suspendDataLoading(for pipeline: ImagePipeline, expectedRequestCount count: Int, _ closure: () -> Void) {
+        let dataLoader = pipeline.configuration.dataLoader as! MockDataLoader
+        dataLoader.isSuspended = true
+        let expectation = self.expectation(description: "registered")
+        expectation.expectedFulfillmentCount = count
+        pipeline.onTaskStarted = { _ in
+            expectation.fulfill()
+        }
         closure()
-        queue.sync {} // Important!
-        configuration.dataLoadingQueue.isSuspended = false
+        wait(for: [expectation], timeout: 5)
+        dataLoader.isSuspended = false
     }
 }