From fc2b28fa531166cb18131c188524f399a40038b0 Mon Sep 17 00:00:00 2001
From: Andrew Patterson <andrew_patterson@intuit.com>
Date: Mon, 23 Dec 2024 12:06:09 -0500
Subject: [PATCH] Only mutate ManagedPlayerViewModel.loadingState on main
 thread

Published properties in SwiftUI models should not be mutated on
background threads.

Refactor `ManagedPlayerViewModel.next` to ensure all paths that update
`loadingState` do so from the main thread.

Add unit tests to verify loadingState changes arrive on the main thread
when next is called.
---
 .../ManagedPlayerViewModel.swift              | 20 +++----
 .../ManagedPlayerViewModelTests.swift         | 58 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/ios/swiftui/Sources/ManagedPlayer/ManagedPlayerViewModel.swift b/ios/swiftui/Sources/ManagedPlayer/ManagedPlayerViewModel.swift
index 19a488231..5d621537a 100644
--- a/ios/swiftui/Sources/ManagedPlayer/ManagedPlayerViewModel.swift
+++ b/ios/swiftui/Sources/ManagedPlayer/ManagedPlayerViewModel.swift
@@ -134,18 +134,16 @@ public class ManagedPlayerViewModel: ObservableObject, NativePlugin {
     }
 
     func next(_ state: CompletedState? = nil) async {
-        DispatchQueue.main.async { [weak self] in
-            self?.loadingState = .loading
-            self?.flow = nil
-        }
-
-        do {
-            let nextState = try await self.manager.next(state)
-            DispatchQueue.main.async { [weak self] in
-                self?.handleNextState(nextState)
+        Task { @MainActor in
+            self.loadingState = .loading
+            self.flow = nil
+
+            do {
+                let nextState = try await self.manager.next(state)
+                self.handleNextState(nextState)
+            } catch {
+                self.loadingState = .failed(error)
             }
-        } catch {
-            self.loadingState = .failed(error)
         }
     }
 
diff --git a/ios/swiftui/Tests/ManagedPlayer/ManagedPlayerViewModelTests.swift b/ios/swiftui/Tests/ManagedPlayer/ManagedPlayerViewModelTests.swift
index 8bc6696fa..f74ac7acd 100644
--- a/ios/swiftui/Tests/ManagedPlayer/ManagedPlayerViewModelTests.swift
+++ b/ios/swiftui/Tests/ManagedPlayer/ManagedPlayerViewModelTests.swift
@@ -84,7 +84,7 @@ class ManagedPlayerViewModelTests: XCTestCase {
         let viewModel = ManagedPlayerViewModel(manager: flowManager, onComplete: {_ in })
 
         await assertPublished(AnyPublisher(viewModel.$loadingState), condition: { state in
-            guard 
+            guard
                 case .failed(let error) = state,
                 let _ = error as? PlayerError
             else { return false }
@@ -281,6 +281,62 @@ class ManagedPlayerViewModelTests: XCTestCase {
             XCTFail("Should have entered idle state")
         }
     }
+
+    func testNextUpdatesLoadingStateOnMainNoFlows() async throws {
+        let model = ManagedPlayerViewModel(manager: ConstantFlowManager([], delay: 1/60), onComplete: {_ in})
+        try await checkNextUpdatesLoadingStateOnMain(model: model, expectedStateChangeCount: 1)
+    }
+
+    func testNextUpdatesLoadingStateOnMainOneFlow() async throws {
+        let model = ManagedPlayerViewModel(manager: ConstantFlowManager(["abc"], delay: 1/60), onComplete: {_ in})
+        try await checkNextUpdatesLoadingStateOnMain(model: model, nextCallCount: 2, expectedStateChangeCount: 4)
+    }
+
+    func testNextUpdatesLoadingStateOnMainMultipleFlows() async throws {
+        let model = ManagedPlayerViewModel(manager: ConstantFlowManager(["abc", "def", "ghi"], delay: 1/60), onComplete: {_ in})
+        try await checkNextUpdatesLoadingStateOnMain(model: model, nextCallCount: 3, expectedStateChangeCount: 6)
+    }
+
+    func testThrowingNextUpdatesLoadingStateOnMain() async throws {
+        let model = ManagedPlayerViewModel(manager: ThrowingFlowManager(), onComplete: {_ in})
+        try await checkNextUpdatesLoadingStateOnMain(model: model, expectedStateChangeCount: 2)
+    }
+
+    private func checkNextUpdatesLoadingStateOnMain(model: ManagedPlayerViewModel, nextCallCount: Int = 1, expectedStateChangeCount: Int) async throws {
+        let expect = (0..<expectedStateChangeCount).map {
+            expectation(description: "\($0)")
+        }
+        var stateChangeCount = 0
+        // the publisher emits immediately, we only care about changes
+        // from the `next` call(s) below, drop the first state received
+        let cancellable = model.$loadingState.dropFirst().sink { state in
+            XCTAssertTrue(Thread.isMainThread)
+            XCTAssertTrue(stateChangeCount < expect.count)
+            stateChangeCount += 1
+            expect[stateChangeCount - 1].fulfill()
+        }
+        await withTaskGroup(of: Void.self) { group in
+            group.addTask {
+                for _ in (0..<nextCallCount) {
+                    await model.next()
+                }
+            }
+            for await _ in group {}
+        }
+        await fulfillment(of: expect)
+        XCTAssertEqual(stateChangeCount, expectedStateChangeCount)
+        cancellable.cancel()
+    }
+
+    private struct ThrowingFlowManager: FlowManager {
+        func next(_: CompletedState?) async throws -> NextState {
+            throw Errors.failed
+        }
+
+        enum Errors: Error {
+            case failed
+        }
+    }
 }
 
 class TerminatingManager: FlowManager {