From d2e9d46da87715b560cb9caf83d11a7e99033d6b Mon Sep 17 00:00:00 2001
From: Jack Newcombe <jack.newcombe@primer.io>
Date: Fri, 19 Apr 2024 22:33:02 +0100
Subject: [PATCH] HUC test fix #2

---
 .../Modules/HUC_TokenizationViewModelTests.swift   | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/Debug App/Tests/Unit Tests/Modules/HUC_TokenizationViewModelTests.swift b/Debug App/Tests/Unit Tests/Modules/HUC_TokenizationViewModelTests.swift
index ca114af536..5371288f0e 100644
--- a/Debug App/Tests/Unit Tests/Modules/HUC_TokenizationViewModelTests.swift	
+++ b/Debug App/Tests/Unit Tests/Modules/HUC_TokenizationViewModelTests.swift	
@@ -22,6 +22,7 @@ final class HUC_TokenizationViewModelTests: XCTestCase {
     private var abortPayment = false
 
     var onCreatePaymentWithData: ((PrimerCheckoutPaymentMethodData) -> Void)?
+    var onUpdateClientSession: ((PrimerClientSession) -> Void?)?
 
     override func tearDown() {
         VaultService.apiClient = nil
@@ -406,12 +407,13 @@ final class HUC_TokenizationViewModelTests: XCTestCase {
         }
 
         if self.abortPayment {
-            onCreatePaymentWithData = { _ in
+
 
                 print(self.eventsCalled)
 
+            onUpdateClientSession = { _ in
+                guard self.eventsCalled.count == 6 else { return }
                 if isSurchargeIncluded {
-                    XCTAssert(self.eventsCalled.count == 6, "6 events should have been called.")
                     XCTAssert(self.eventsCalled[0] == "primerHeadlessUniversalCheckoutPreparationDidStart", "'\(self.eventsCalled[0])' called instead if 'primerHeadlessUniversalCheckoutPreparationDidStart'.")
                     XCTAssert(self.eventsCalled[1] == "primerHeadlessUniversalCheckoutClientSessionWillUpdate", "'\(self.eventsCalled[1])' called instead if 'primerHeadlessUniversalCheckoutClientSessionWillUpdate'.")
                     XCTAssert(self.eventsCalled[2] == "primerHeadlessUniversalCheckoutClientSessionDidUpdate", "'\(self.eventsCalled[2])' called instead if 'primerHeadlessUniversalCheckoutClientSessionDidUpdate'.")
@@ -419,7 +421,12 @@ final class HUC_TokenizationViewModelTests: XCTestCase {
                     XCTAssert(self.eventsCalled[4] == "primerHeadlessUniversalCheckoutClientSessionWillUpdate", "'\(self.eventsCalled[4])' called instead if 'primerHeadlessUniversalCheckoutClientSessionWillUpdate'.")
                     XCTAssert(self.eventsCalled[5] == "primerHeadlessUniversalCheckoutClientSessionDidUpdate", "'\(self.eventsCalled[5])' called instead if 'primerHeadlessUniversalCheckoutClientSessionDidUpdate'.")
 
-                } else {
+                }
+                expectation.fulfill()
+            }
+
+            onCreatePaymentWithData = { _ in
+                if !isSurchargeIncluded {
                     XCTAssert(self.eventsCalled.count == 2, "2 events should have been called but got \(self.eventsCalled.count).")
                     XCTAssert(self.eventsCalled[0] == "primerHeadlessUniversalCheckoutPreparationDidStart", "'\(self.eventsCalled[0])' called instead if 'primerHeadlessUniversalCheckoutPreparationDidStart'.")
                     XCTAssert(self.eventsCalled[1] == "primerHeadlessUniversalCheckoutWillCreatePaymentWithData", "'\(self.eventsCalled[1])' called instead if 'primerHeadlessUniversalCheckoutWillCreatePaymentWithData'.")
@@ -478,6 +485,7 @@ extension HUC_TokenizationViewModelTests: PrimerHeadlessUniversalCheckoutDelegat
 
     func primerHeadlessUniversalCheckoutDidUpdateClientSession(_ clientSession: PrimerClientSession) {
         eventsCalled.append("primerHeadlessUniversalCheckoutClientSessionDidUpdate")
+        onUpdateClientSession?(clientSession)
     }
 
     func primerHeadlessUniversalCheckoutDidReceiveAdditionalInfo(_ additionalInfo: PrimerCheckoutAdditionalInfo?) {