From a19c933c37007aedf9d94b16958da229493bc4c1 Mon Sep 17 00:00:00 2001
From: Gus Cairo <gus.cairo@apple.com>
Date: Tue, 15 Apr 2025 15:08:35 +0100
Subject: [PATCH] PR changes

---
 .../NIOTSBootstrapTests.swift                        | 12 ++++++++----
 .../NIOTSDatagramConnectionChannelTests.swift        | 10 +++++++---
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift b/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
index 969597a..4a00dd4 100644
--- a/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
+++ b/Tests/NIOTransportServicesTests/NIOTSBootstrapTests.swift
@@ -18,7 +18,7 @@ import XCTest
 import Network
 import NIOCore
 import NIOEmbedded
-@testable import NIOTransportServices
+import NIOTransportServices
 import NIOConcurrencyHelpers
 import Foundation
 
@@ -384,10 +384,16 @@ final class NIOTSBootstrapTests: XCTestCase {
 
             func channelActive(context: ChannelHandlerContext) {
                 self.connectionPromise.succeed()
+                context.pipeline.fireChannelActive()
             }
         }
 
-        let group = NIOTSEventLoopGroup(loopCount: 1)
+        let group = NIOTSEventLoopGroup()
+        defer {
+            Task {
+                try await group.shutdownGracefully()
+            }
+        }
 
         let configuratorListenerCounter = NIOLockedValueBox(0)
         let configuratorConnectionCounter = NIOLockedValueBox(0)
@@ -422,8 +428,6 @@ final class NIOTSBootstrapTests: XCTestCase {
 
         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 39e1c99..eca509b 100644
--- a/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift
+++ b/Tests/NIOTransportServicesTests/NIOTSDatagramConnectionChannelTests.swift
@@ -245,10 +245,16 @@ final class NIOTSDatagramConnectionChannelTests: XCTestCase {
 
             func channelActive(context: ChannelHandlerContext) {
                 self.connectionPromise.succeed()
+                context.fireChannelActive()
             }
         }
 
-        let group = NIOTSEventLoopGroup(loopCount: 1)
+        let group = NIOTSEventLoopGroup()
+        defer {
+            Task {
+                try await group.shutdownGracefully()
+            }
+        }
 
         let configuratorListenerCounter = NIOLockedValueBox(0)
         let configuratorConnectionCounter = NIOLockedValueBox(0)
@@ -288,8 +294,6 @@ final class NIOTSDatagramConnectionChannelTests: XCTestCase {
 
         XCTAssertEqual(2, configuratorListenerCounter.withLockedValue { $0 })
         XCTAssertEqual(1, configuratorConnectionCounter.withLockedValue { $0 })
-
-        try await group.shutdownGracefully()
     }
 
     func testCanExtractTheConnection() throws {