Fix broken upgrade

apple · Dec 13, 2024 · 6fbecdd · 6fbecdd
1 parent aa646a8
commit 6fbecdd
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift b/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift
@@ -684,8 +684,6 @@ final class TypedWebSocketClientEndToEndTests: WebSocketClientEndToEndTests {
     }
 
     override fileprivate func runSuccessfulUpgrade() throws -> (EmbeddedChannel, WebSocketRecorderHandler) {
-        let handler = WebSocketRecorderHandler()
-
         let basicUpgrader = NIOTypedWebSocketClientUpgrader(
             requestKey: "OfS0wDaT5NoxF2gqm7Zj2YtetzM=",
             upgradePipelineHandler: { (channel: Channel, _: HTTPResponseHead) in
@@ -716,6 +714,10 @@ final class TypedWebSocketClientEndToEndTests: WebSocketClientEndToEndTests {
 
         try upgradeResult.wait()
 
+        // Ok, now grab the handler. We can do this with sync operations, because this is an
+        // EmbeddedChannel.
+        let handler = try clientChannel.pipeline.syncOperations.handler(type: WebSocketRecorderHandler.self)
+
         return (clientChannel, handler)
     }
 }