Loading animation only works for first image until fetched

Question

I am trying to hide the loading process of different elements in my app. My loading wheel only works until the first image appears. However my app is fetching two images, but due to how the site is loading everything its not synchronized. How would I update it so the spinning wheel stays on until everything is fetched?

import SwiftUI
import WebKit

struct ContentView: View {
    @StateObject var webViewStore = WebViewStore()

    var body: some View {
        VStack(alignment: .leading) {
            contentSection
            Divider()
            urlInputSection
        }
        .padding()
        .background(Color(red: 0.15, green: 0.15, blue: 0.15))
    }

    private var contentSection: some View {
        Group {
            if webViewStore.isLoading {
                loadingView
            } else {
                VStack {
                    HStack(alignment: .top) {
                        if let uiImage = webViewStore.fetchedUIImage {
                            Image(uiImage: uiImage)
                                .resizable()
                                .aspectRatio(contentMode: .fit)
                                .frame(width: 100, height: 100)
                                .cornerRadius(5)
                        }

                        VStack(alignment: .leading) {
                            userProfileSection
                            descriptionSection
                            Spacer()
                        }
                        .padding(.leading, 10)
                    }

                    Spacer()
                }
            }
        }
    }

    private var userProfileSection: some View {
        HStack {
            if let userProfileImage = webViewStore.fetchedProfileUIImage {
                Image(uiImage: userProfileImage)
                    .resizable()
                    .aspectRatio(contentMode: .fit)
                    .frame(width: 100, height: 100)
            }

            if let username = webViewStore.fetchedUsername {
                Text(username)
                    .font(.caption)
                    .foregroundColor(.white)
            }
        }
    }

    private var descriptionSection: some View {
        Group {
            if let description = webViewStore.fetchedDescription {
                Text(description)
                    .font(.caption)
                    .foregroundColor(.white)
                    .lineLimit(4)
                    .padding(.top, 5)
            }
        }
    }

    private var urlInputSection: some View {
        HStack {
            Button(action: {
                webViewStore.loadURL(url: URL(string: "https://www.instagram.com/p/CxGEZGWrH-N/")!)
            }) {
                Text("Load")
                    .foregroundColor(.white)
                    .padding(.horizontal, 16)
                    .padding(.vertical, 8)
                    .background(Color.blue)
                    .cornerRadius(8)
            }
            .frame(height: 36)
        }
    }

    private var loadingView: some View {
        VStack {
            ProgressView("Loading...")
                .padding()
        }
        .frame(maxWidth: .infinity, maxHeight: .infinity)
        .background(Color(red: 0.15, green: 0.15, blue: 0.15))
    }
}

class WebViewStore: NSObject, ObservableObject, WKNavigationDelegate {
    @Published var fetchedUIImage: UIImage?
    @Published var fetchedDescription: String?
    @Published var fetchedProfileUIImage: UIImage?
    @Published var fetchedUsername: String?
    @Published var isLoading: Bool = false
    
    private var webView: WKWebView

    override init() {
        let webConfiguration = WKWebViewConfiguration()
        self.webView = WKWebView(frame: .zero, configuration: webConfiguration)
        super.init()
        self.webView.navigationDelegate = self
    }

    func loadURL(url: URL) {
        isLoading = true
        let request = URLRequest(url: url)
        webView.load(request)
    }

    func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
        fetchContent()
        attemptFetchingDescription(retryCount: 3)
    }
    
    private func fetchContent() {
        let fetchImageScript = """
        (function() {
            var imageElement = document.querySelector('meta[property="og:image"]');
            var result = {
                imageURL: null,
            };
            if (imageElement) {
                result.imageURL = imageElement.getAttribute('content');
            }
            return result;
        })();
        """

        webView.evaluateJavaScript(fetchImageScript) { (fetchResult, fetchError) in
            if let result = fetchResult as? [String: String], let imageURLString = result["imageURL"], let validURL = URL(string: imageURLString) {
                self.loadImage(from: validURL, assignTo: \.fetchedUIImage)
            }
        }
    }
    
    private func loadImage(from url: URL, assignTo property: ReferenceWritableKeyPath<WebViewStore, UIImage?>) {
        URLSession.shared.dataTask(with: url) { (data, response, error) in
            if let imageData = data, let uiImage = UIImage(data: imageData) {
                DispatchQueue.main.async {
                    self[keyPath: property] = uiImage
                    self.isLoading = false
                }
            }
        }.resume()
    }

    private func attemptFetchingDescription(retryCount: Int) {
        if retryCount <= 0 {
            isLoading = false
            return
        }

        let checkScript = "document.body.innerHTML.includes('Instagram')"

        DispatchQueue.main.async {
            self.webView.evaluateJavaScript(checkScript) { (result, error) in
                if let isPresent = result as? Bool, isPresent {
                    self.fetchImageURL(retryCount: retryCount)
                } else {
                    DispatchQueue.global().asyncAfter(deadline: .now() + 1.0) {
                        self.attemptFetchingDescription(retryCount: retryCount - 1)
                    }
                }
            }
        }
    }

    private func fetchImageURL(retryCount: Int) {
        let fetchImageScript = "var img = document.querySelector('header img.xpdipgo.x972fbf'); img ? img.getAttribute('src') : 'Not found';"
        DispatchQueue.main.async {
            self.webView.evaluateJavaScript(fetchImageScript) { (fetchResult, fetchError) in
                if let link = fetchResult as? String, !link.contains("Not found"), let validURL = URL(string: link) {
                    self.loadImage(from: validURL, assignTo: \.fetchedProfileUIImage)
                } else {
                    DispatchQueue.global().asyncAfter(deadline: .now() + 2.0) {
                        self.attemptFetchingDescription(retryCount: retryCount - 1)
                    }
                }
            }
        }
    }
}

Answer 1

You can achieve this using DispatchGroup.

1. Add a property called just under your class definition: private let fetchGroup = DispatchGroup()

2. Modify loadImage method: Enter the fetchGroup when you start the loading process and leave the fetchGroup after the image has been loaded:

   private func loadImage(from url: URL, assignTo property: ReferenceWritableKeyPath<WebViewStore, UIImage?>) {
   fetchGroup.enter()
   URLSession.shared.dataTask(with: url) { (data, response, error) in
       if let imageData = data, let uiImage = UIImage(data: imageData) {
           DispatchQueue.main.async {
               self[keyPath: property] = uiImage
           }
       }
       self.fetchGroup.leave()
   }.resume()
   

3. After the images are loaded, wait for the fetchGroup to complete, and then set isLoading to false:

   func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
   fetchContent()
   attemptFetchingDescription(retryCount: 3)
   
   fetchGroup.notify(queue: .main) {
       self.isLoading = false
   }