From 85cb411b3bf5e8017621fae1ea4d4116b2407ab8 Mon Sep 17 00:00:00 2001
From: Artur Guseinov <arthur.guseinov@gmail.com>
Date: Tue, 27 Sep 2022 14:57:25 +0600
Subject: [PATCH] Combine timeout for dispatcher logic

---
 Example/IntegrationTests/Auth/AuthTests.swift | 17 ------------
 Example/IntegrationTests/Chat/ChatTests.swift | 21 ---------------
 .../Sign/SignClientTests.swift                | 15 -----------
 Sources/WalletConnectRelay/Dispatching.swift  | 27 ++++++++++++-------
 4 files changed, 17 insertions(+), 63 deletions(-)

diff --git a/Example/IntegrationTests/Auth/AuthTests.swift b/Example/IntegrationTests/Auth/AuthTests.swift
index 6abe13983..4ee2bf33d 100644
--- a/Example/IntegrationTests/Auth/AuthTests.swift
+++ b/Example/IntegrationTests/Auth/AuthTests.swift
@@ -16,23 +16,6 @@ final class AuthTests: XCTestCase {
         app = makeClient(prefix: "👻 App")
         let walletAccount = Account(chainIdentifier: "eip155:1", address: "0x724d0D2DaD3fbB0C168f947B87Fa5DBe36F1A8bf")!
         wallet = makeClient(prefix: "🤑 Wallet", account: walletAccount)
-
-        let expectation = expectation(description: "Wait Clients Connected")
-        expectation.expectedFulfillmentCount = 2
-
-        app.socketConnectionStatusPublisher.sink { status in
-            if status == .connected {
-                expectation.fulfill()
-            }
-        }.store(in: &publishers)
-
-        wallet.socketConnectionStatusPublisher.sink { status in
-            if status == .connected {
-                expectation.fulfill()
-            }
-        }.store(in: &publishers)
-
-        wait(for: [expectation], timeout: 5)
     }
 
     func makeClient(prefix: String, account: Account? = nil) -> AuthClient {
diff --git a/Example/IntegrationTests/Chat/ChatTests.swift b/Example/IntegrationTests/Chat/ChatTests.swift
index 191ea398f..c4863b8af 100644
--- a/Example/IntegrationTests/Chat/ChatTests.swift
+++ b/Example/IntegrationTests/Chat/ChatTests.swift
@@ -16,27 +16,6 @@ final class ChatTests: XCTestCase {
         registry = KeyValueRegistry()
         invitee = makeClient(prefix: "🦖 Registered")
         inviter = makeClient(prefix: "🍄 Inviter")
-
-        waitClientsConnected()
-    }
-
-    private func waitClientsConnected() {
-        let expectation = expectation(description: "Wait Clients Connected")
-        expectation.expectedFulfillmentCount = 2
-
-        invitee.socketConnectionStatusPublisher.sink { status in
-            if status == .connected {
-                expectation.fulfill()
-            }
-        }.store(in: &publishers)
-
-        inviter.socketConnectionStatusPublisher.sink { status in
-            if status == .connected {
-                expectation.fulfill()
-            }
-        }.store(in: &publishers)
-
-        wait(for: [expectation], timeout: 5)
     }
 
     func makeClient(prefix: String) -> ChatClient {
diff --git a/Example/IntegrationTests/Sign/SignClientTests.swift b/Example/IntegrationTests/Sign/SignClientTests.swift
index a62321a05..20bd91504 100644
--- a/Example/IntegrationTests/Sign/SignClientTests.swift
+++ b/Example/IntegrationTests/Sign/SignClientTests.swift
@@ -37,24 +37,9 @@ final class SignClientTests: XCTestCase {
         return ClientDelegate(client: client)
     }
 
-    private func listenForConnection() async {
-        let group = DispatchGroup()
-        group.enter()
-        dapp.onConnected = {
-            group.leave()
-        }
-        group.enter()
-        wallet.onConnected = {
-            group.leave()
-        }
-        group.wait()
-        return
-    }
-
     override func setUp() async throws {
         dapp = Self.makeClientDelegate(name: "🍏P")
         wallet = Self.makeClientDelegate(name: "🍎R")
-        await listenForConnection()
     }
 
     override func tearDown() {
diff --git a/Sources/WalletConnectRelay/Dispatching.swift b/Sources/WalletConnectRelay/Dispatching.swift
index 6df8a56f4..cd08efe12 100644
--- a/Sources/WalletConnectRelay/Dispatching.swift
+++ b/Sources/WalletConnectRelay/Dispatching.swift
@@ -26,6 +26,8 @@ final class Dispatcher: NSObject, Dispatching {
         socketConnectionStatusPublisherSubject.eraseToAnyPublisher()
     }
 
+    private let concurrentQueue = DispatchQueue(label: "com.walletconnect.sdk.dispatcher", attributes: .concurrent)
+
     init(socket: WebSocketConnecting,
          socketConnectionHandler: SocketConnectionHandler,
          logger: ConsoleLogging) {
@@ -53,16 +55,21 @@ final class Dispatcher: NSObject, Dispatching {
         }
 
         var cancellable: AnyCancellable?
-        cancellable = socketConnectionStatusPublisher.sink { [unowned self] status in
-            guard status == .connected else { return }
-            cancellable?.cancel()
-            send(string, completion: completion)
-        }
-
-        DispatchQueue.global().asyncAfter(deadline: .now() + .seconds(defaultTimeout)) {
-            completion(NetworkError.webSocketNotConnected)
-            cancellable?.cancel()
-        }
+        cancellable = socketConnectionStatusPublisher
+            .filter { $0 == .connected }
+            .setFailureType(to: NetworkError.self)
+            .timeout(.seconds(defaultTimeout), scheduler: concurrentQueue, customError: { .webSocketNotConnected })
+            .sink(receiveCompletion: { result in
+                switch result {
+                case .failure(let error):
+                    cancellable?.cancel()
+                    completion(error)
+                case .finished: break
+                }
+            }, receiveValue: { [unowned self] status in
+                cancellable?.cancel()
+                send(string, completion: completion)
+            })
     }
 
     func protectedSend(_ string: String) async throws {