From 19b43c48e06404e0d489583e5106cad01929a5ab Mon Sep 17 00:00:00 2001
From: Hesham Salman <iron-ham@github.com>
Date: Tue, 20 Aug 2024 13:58:53 -0400
Subject: [PATCH] Fix flaky OffsetPagination test

---
 Tests/ApolloPaginationTests/OffsetTests.swift | 37 +++++++++++--------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/Tests/ApolloPaginationTests/OffsetTests.swift b/Tests/ApolloPaginationTests/OffsetTests.swift
index e917ca6fd..de069068b 100644
--- a/Tests/ApolloPaginationTests/OffsetTests.swift
+++ b/Tests/ApolloPaginationTests/OffsetTests.swift
@@ -89,8 +89,13 @@ final class OffsetTests: XCTestCase {
     }
 
     let pager = await createPager()
-    var results: [ViewModel]?
-    let cancellable = pager.map { value in
+
+    let fetchExpectation = expectation(description: "Initial Fetch")
+    fetchExpectation.assertForOverFulfill = false
+    let subscriptionExpectation = expectation(description: "Subscription")
+    subscriptionExpectation.expectedFulfillmentCount = 2
+    var expectedViewModels: [ViewModel] = []
+    let subscriber = pager.compactMap { value in
       switch value {
       case .success(let output):
         let friends = output.allData.flatMap { data in
@@ -98,29 +103,29 @@ final class OffsetTests: XCTestCase {
             ViewModel(name: friend.name)
           }
         }
-        return Result<[ViewModel], any Error>.success(friends)
+        return friends
       case .failure(let error):
-        return .failure(error)
-      }
-    }.sink { result in
-      switch result {
-      case .success((let viewModels)):
-        results = viewModels
-      default:
-        XCTFail("Failed to get view models from pager.")
+        XCTFail(error.localizedDescription)
+        return nil
       }
+    }.sink { viewModels in
+      expectedViewModels = viewModels
+      fetchExpectation.fulfill()
+      subscriptionExpectation.fulfill()
     }
 
     await fetchFirstPage(pager: pager)
-    XCTAssertEqual(results?.count, 2)
-    XCTAssertEqual(results?.map(\.name), ["Luke Skywalker", "Han Solo"])
+    await fulfillment(of: [fetchExpectation], timeout: 1)
+    XCTAssertEqual(expectedViewModels.count, 2)
+    XCTAssertEqual(expectedViewModels.map(\.name), ["Luke Skywalker", "Han Solo"])
     let canLoadNext = await pager.canLoadNext
     XCTAssertTrue(canLoadNext)
 
     try await fetchSecondPage(pager: pager)
-    XCTAssertEqual(results?.count, 3)
-    XCTAssertEqual(results?.map(\.name), ["Luke Skywalker", "Han Solo", "Leia Organa"])
-    cancellable.cancel()
+    await fulfillment(of: [subscriptionExpectation], timeout: 1)
+    XCTAssertEqual(expectedViewModels.count, 3)
+    XCTAssertEqual(expectedViewModels.map(\.name), ["Luke Skywalker", "Han Solo", "Leia Organa"])
+    subscriber.cancel()
   }
 
 }