Fix multi callback from stateUpdateHandler issue

Author: Kyle-Ye

Sources/OpenGraphShims/DebugClient.swift

@@ -4,42 +4,51 @@
 
 #if canImport(Darwin)
 public import Foundation
-import Network
+public import Network
 
 #if canImport(OpenGraphCxx_Private)
 public import OpenGraphCxx_Private.DebugServer
 #endif
 
+public struct ConnectionUpdates: AsyncSequence {
+    public typealias Element = NWConnection.State
+    
+    private let stream: AsyncStream<NWConnection.State>
+    
+    fileprivate init(stream: AsyncStream<NWConnection.State>) {
+        self.stream = stream
+    }
+    
+    public func makeAsyncIterator() -> AsyncStream<NWConnection.State>.AsyncIterator {
+        stream.makeAsyncIterator()
+    }
+}
+
 @_spi(Debug)
 public final class DebugClient {
     private var connection: NWConnection?
-    private let queue = DispatchQueue(label: "opengraph.debugserver.client.queue")
+    private let queue = DispatchQueue(label: "org.openswiftuiproject.opengraph.debugclient")
 
-    public func connect(to url: URL) async throws {
+    public func connect(to url: URL) -> ConnectionUpdates {
         guard let host = url.host, let port = url.port else {
-            throw ClientError.invalidURL
+            return ConnectionUpdates(stream: AsyncStream { continuation in
+                continuation.yield(.failed(NWError.posix(.EINVAL)))
+                continuation.finish()
+            })
         }
-
         let nwHost = NWEndpoint.Host(host)
         let nwPort = NWEndpoint.Port(integerLiteral: UInt16(port))
-
         connection = NWConnection(host: nwHost, port: nwPort, using: .tcp)
-
-        return try await withCheckedThrowingContinuation { continuation in
+        let stream = AsyncStream<NWConnection.State> { continuation in
             connection?.stateUpdateHandler = { state in
-                switch state {
-                case .ready:
-                    continuation.resume()
-                case let .failed(error):
-                    continuation.resume(throwing: error)
-                case .cancelled:
-                    continuation.resume(throwing: ClientError.connectionCancelled)
-                default:
-                    break
+                continuation.yield(state)
+                if case .cancelled = state {
+                    continuation.finish()
                 }
             }
             connection?.start(queue: queue)
         }
+        return ConnectionUpdates(stream: stream)
     }
 
     public func sendMessage(token: UInt32, data: Data) async throws {

Tests/OpenGraphCxxTests/DebugServer/DebugServerTests.swift

@@ -24,26 +24,37 @@ struct DebugServerTests {
 
     @Test
     func commandTest() async throws {
-        let debugServer = OG.DebugServer([.valid])
+        var debugServer = OG.DebugServer([.valid])
         let cfURL = debugServer.copy_url()
         let url = try #require(cfURL) as URL
         let components = try #require(URLComponents(url: url, resolvingAgainstBaseURL: false))
         let token = try #require(components.queryItems?.first { $0.name == "token" }?.value.flatMap { UInt32($0) })
         debugServer.run(1)
         let client = DebugClient()
-        try await client.connect(to: url)
-
-        for command in Command.allCases {
-            if command == .graphDescription {
-                continue
+        let updates = client.connect(to: url)
+        try await confirmation { confirm in
+            for try await update in updates {
+                switch update {
+                case .ready:
+                    confirm()
+                    for command in Command.allCases {
+                        if command == .graphDescription {
+                            continue
+                        }
+                        try await client.sendMessage(
+                            token: token,
+                            data: data(for: command)
+                        )
+                        let (_, responseData) = try await client.receiveMessage()
+                        let response = try #require(String(data: responseData, encoding: .utf8))
+                        #expect(response == command.rawValue)
+                    }
+                    debugServer.shutdown()
+                    // TODO: The shutdown should close the connection, but it does not for OGDebugServer currently.
+                default:
+                    break
+                }
             }
-            try await client.sendMessage(
-                token: token,
-                data: data(for: command)
-            )
-            let (_, responseData) = try await client.receiveMessage()
-            let response = try #require(String(data: responseData, encoding: .utf8))
-            #expect(response == command.rawValue)
         }
     }
 }