From 3a68b631837f96dd66a32fb9fa3e4d569194c483 Mon Sep 17 00:00:00 2001
From: Guille Gonzalez <guillermo.gonzalez@me.com>
Date: Thu, 17 Jun 2021 11:57:51 +0200
Subject: [PATCH] Update cached images immediately

---
 .../Core/NetworkImageLoader.swift             | 65 +++++++++++++------
 .../NetworkImage/Core/NetworkImageStore.swift | 16 +++--
 .../NetworkImageLoaderTests.swift             |  1 +
 3 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/Sources/NetworkImage/Core/NetworkImageLoader.swift b/Sources/NetworkImage/Core/NetworkImageLoader.swift
index 944b748..9e5bd5b 100644
--- a/Sources/NetworkImage/Core/NetworkImageLoader.swift
+++ b/Sources/NetworkImage/Core/NetworkImageLoader.swift
@@ -7,6 +7,7 @@
     @available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
     public struct NetworkImageLoader {
         private let _image: (URL) -> AnyPublisher<OSImage, Error>
+        private let _cachedImage: (URL) -> OSImage?
 
         /// Creates an image loader.
         /// - Parameters:
@@ -17,38 +18,52 @@
         }
 
         init(urlLoader: URLLoader, imageCache: NetworkImageCache) {
-            self.init { url in
-                if let image = imageCache.image(for: url) {
-                    return Just(image)
-                        .setFailureType(to: Error.self)
-                        .eraseToAnyPublisher()
-                } else {
-                    return urlLoader.dataTaskPublisher(for: url)
-                        .tryMap { data, response in
-                            if let httpResponse = response as? HTTPURLResponse {
-                                guard 200 ..< 300 ~= httpResponse.statusCode else {
-                                    throw NetworkImageError.badStatus(httpResponse.statusCode)
+            self.init(
+                image: { url in
+                    if let image = imageCache.image(for: url) {
+                        return Just(image)
+                            .setFailureType(to: Error.self)
+                            .eraseToAnyPublisher()
+                    } else {
+                        return urlLoader.dataTaskPublisher(for: url)
+                            .tryMap { data, response in
+                                if let httpResponse = response as? HTTPURLResponse {
+                                    guard 200 ..< 300 ~= httpResponse.statusCode else {
+                                        throw NetworkImageError.badStatus(httpResponse.statusCode)
+                                    }
                                 }
-                            }
 
-                            return try decodeImage(from: data)
-                        }
-                        .handleEvents(receiveOutput: { image in
-                            imageCache.setImage(image, for: url)
-                        })
-                        .eraseToAnyPublisher()
+                                return try decodeImage(from: data)
+                            }
+                            .handleEvents(receiveOutput: { image in
+                                imageCache.setImage(image, for: url)
+                            })
+                            .eraseToAnyPublisher()
+                    }
+                },
+                cachedImage: { url in
+                    imageCache.image(for: url)
                 }
-            }
+            )
         }
 
-        init(image: @escaping (URL) -> AnyPublisher<OSImage, Error>) {
+        init(
+            image: @escaping (URL) -> AnyPublisher<OSImage, Error>,
+            cachedImage: @escaping (URL) -> OSImage?
+        ) {
             _image = image
+            _cachedImage = cachedImage
         }
 
         /// Returns a publisher that loads an image for a given URL.
         public func image(for url: URL) -> AnyPublisher<OSImage, Error> {
             _image(url)
         }
+
+        /// Returns the cached image for a given URL if there is any.
+        public func cachedImage(for url: URL) -> OSImage? {
+            _cachedImage(url)
+        }
     }
 
     @available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
@@ -73,13 +88,19 @@
                     }
 
                     return response.eraseToAnyPublisher()
+                } cachedImage: { _ in
+                    nil
                 }
             }
 
             static func mock<P>(
                 response: P
             ) -> Self where P: Publisher, P.Output == OSImage, P.Failure == Error {
-                Self { _ in response.eraseToAnyPublisher() }
+                Self { _ in
+                    response.eraseToAnyPublisher()
+                } cachedImage: { _ in
+                    nil
+                }
             }
 
             static var failing: Self {
@@ -88,6 +109,8 @@
                     return Just(OSImage())
                         .setFailureType(to: Error.self)
                         .eraseToAnyPublisher()
+                } cachedImage: { _ in
+                    nil
                 }
             }
         }
diff --git a/Sources/NetworkImage/Core/NetworkImageStore.swift b/Sources/NetworkImage/Core/NetworkImageStore.swift
index 5e3ad84..e75b752 100644
--- a/Sources/NetworkImage/Core/NetworkImageStore.swift
+++ b/Sources/NetworkImage/Core/NetworkImageStore.swift
@@ -21,13 +21,17 @@
 
         init(url: URL?, environment: NetworkImageEnvironment) {
             if let url = url {
-                state = .placeholder
+                if let image = environment.imageLoader.cachedImage(for: url) {
+                    state = .image(image)
+                } else {
+                    state = .placeholder
 
-                environment.imageLoader.image(for: url)
-                    .map { .image($0) }
-                    .replaceError(with: .fallback)
-                    .receive(on: environment.mainQueue)
-                    .assign(to: &$state)
+                    environment.imageLoader.image(for: url)
+                        .map { .image($0) }
+                        .replaceError(with: .fallback)
+                        .receive(on: environment.mainQueue)
+                        .assign(to: &$state)
+                }
             } else {
                 state = .fallback
             }
diff --git a/Tests/NetworkImageTests/NetworkImageLoaderTests.swift b/Tests/NetworkImageTests/NetworkImageLoaderTests.swift
index cd2e689..47ff685 100644
--- a/Tests/NetworkImageTests/NetworkImageLoaderTests.swift
+++ b/Tests/NetworkImageTests/NetworkImageLoaderTests.swift
@@ -46,6 +46,7 @@
             // then
             let unwrappedResult = try XCTUnwrap(result)
             XCTAssertTrue(unwrappedResult.isEqual(imageCache.image(for: Fixtures.anyImageURL)))
+            XCTAssertTrue(unwrappedResult.isEqual(imageLoader.cachedImage(for: Fixtures.anyImageURL)))
         }
 
         func testImageReturnsCachedImageIfAvailable() throws {