From c90c6432fd5abb3ad3b771a2782a3d34c0fe730f Mon Sep 17 00:00:00 2001
From: Gus Cairo <gus.cairo@apple.com>
Date: Mon, 14 Apr 2025 14:18:35 +0100
Subject: [PATCH] Add tests

---
 .../NIOTSBootstrapTests.swift                 | 30 +++++++++-------
 .../NIOTSDatagramConnectionChannelTests.swift | 34 ++++++++++++-------
 2 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift b/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
index c6242cc..3d4264e 100644
--- a/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
+++ b/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
@@ -18,7 +18,7 @@ import XCTest
 import Network
 import NIOCore
 import NIOEmbedded
-import NIOTransportServices
+@testable import NIOTransportServices
 import NIOConcurrencyHelpers
 import Foundation
 
@@ -372,34 +372,40 @@ final class NIOTSBootstrapTests: XCTestCase {
         XCTAssertEqual(try connectionChannel.getOption(NIOTSChannelOptions.multipathServiceType).wait(), .handover)
     }
 
-    func testNWParametersConfigurator() throws {
+    func testNWParametersConfigurator() async throws {
         let group = NIOTSEventLoopGroup()
-        defer {
-            try! group.syncShutdownGracefully()
-        }
 
         let configuratorListenerCounter = NIOLockedValueBox(0)
         let configuratorConnectionCounter = NIOLockedValueBox(0)
 
-        let listenerChannel = try NIOTSListenerBootstrap(group: group)
+        let listenerChannel = try await NIOTSListenerBootstrap(group: group)
             .configureNWParameters { _ in
                 configuratorListenerCounter.withLockedValue { $0 += 1 }
             }
             .bind(host: "localhost", port: 0)
-            .wait()
+            .get()
 
-        let connectionChannel: Channel = try NIOTSConnectionBootstrap(group: group)
+        let connectionChannel: Channel = try await NIOTSConnectionBootstrap(group: group)
             .configureNWParameters { _ in
                 configuratorConnectionCounter.withLockedValue { $0 += 1 }
             }
             .connect(to: listenerChannel.localAddress!)
-            .wait()
+            .get()
 
-        try listenerChannel.close().wait()
-        try connectionChannel.close().wait()
+        // Need to wait for the connection channel to be activated. We cannot wait for the connect
+        // promise (`StateManagedNWConnectionChannel/connectPromise`) from here, because it is
+        // nulled out as soon as it's completed, and the listener channel doesn't wait on it either.
+        // If we don't wait, we can race with the connection channel's parameter configurator
+        // closure being run, and the assertion below would fail.
+        try await Task.sleep(for: .milliseconds(100))
 
-        XCTAssertEqual(1, configuratorListenerCounter.withLockedValue { $0 })
+        try await listenerChannel.close().get()
+        try await connectionChannel.close().get()
+
+        XCTAssertEqual(2, configuratorListenerCounter.withLockedValue { $0 })
         XCTAssertEqual(1, configuratorConnectionCounter.withLockedValue { $0 })
+
+        try await group.shutdownGracefully()
     }
 }
 
diff --git a/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift b/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift
index b9f8b2f..8efd327 100644
--- a/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift
+++ b/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift
@@ -16,7 +16,7 @@
 import XCTest
 import Network
 import NIOCore
-import NIOTransportServices
+@testable import NIOTransportServices
 import Foundation
 import NIOConcurrencyHelpers
 
@@ -233,34 +233,44 @@ final class NIOTSDatagramConnectionChannelTests: XCTestCase {
         XCTAssertNoThrow(try connection.close().wait())
     }
 
-    func testNWParametersConfigurator() throws {
+    func testNWParametersConfigurator() async throws {
         let group = NIOTSEventLoopGroup()
-        defer {
-            try! group.syncShutdownGracefully()
-        }
 
         let configuratorListenerCounter = NIOLockedValueBox(0)
         let configuratorConnectionCounter = NIOLockedValueBox(0)
 
-        let listenerChannel = try NIOTSDatagramListenerBootstrap(group: group)
+        let listenerChannel = try await NIOTSDatagramListenerBootstrap(group: group)
+//            .childChannelInitializer { connectionChannel in
+//                print((connectionChannel as! NIOTSDatagramChannel).connectPromise)
+//                return connectionChannel.eventLoop.makeSucceededFuture(())
+//            }
             .configureNWParameters { _ in
                 configuratorListenerCounter.withLockedValue { $0 += 1 }
             }
             .bind(host: "localhost", port: 0)
-            .wait()
+            .get()
 
-        let connectionChannel: Channel = try NIOTSDatagramBootstrap(group: group)
+        let connectionChannel: Channel = try await NIOTSDatagramBootstrap(group: group)
             .configureNWParameters { _ in
                 configuratorConnectionCounter.withLockedValue { $0 += 1 }
             }
             .connect(to: listenerChannel.localAddress!)
-            .wait()
+            .get()
 
-        try listenerChannel.close().wait()
-        try connectionChannel.close().wait()
+        // Need to write something so the server can activate the connection channel: this is UDP,
+        // so there is no handshaking that happens and thus the server cannot know that the
+        // connection has been established and the channel can be activated.
+        try await connectionChannel.writeAndFlush(ByteBuffer(bytes: [42]))
 
-        XCTAssertEqual(1, configuratorListenerCounter.withLockedValue { $0 })
+        try await Task.sleep(for: .milliseconds(100))
+
+        try await listenerChannel.close().get()
+        try await connectionChannel.close().get()
+
+        XCTAssertEqual(2, configuratorListenerCounter.withLockedValue { $0 })
         XCTAssertEqual(1, configuratorConnectionCounter.withLockedValue { $0 })
+
+        try await group.shutdownGracefully()
     }
 
     func testCanExtractTheConnection() throws {