From 1352aa2dee5d3c1283b53823e8461587554c1e87 Mon Sep 17 00:00:00 2001
From: Adam Fowler <adamfowler71@gmail.com>
Date: Mon, 17 Jun 2024 10:58:00 +0100
Subject: [PATCH] RequestContext with associatedtype Source (#64)

* Update for context source

* Use context-source branch until it is merged

* Get working with collapsed RequestContext

* Fixes for name changes

* Use Humminbird 2.0.0-beta.7
---
 Package.swift                                        |  2 +-
 Sources/HummingbirdWebSocket/WebSocketRouter.swift   | 12 ++++++------
 Tests/HummingbirdWebSocketTests/WebSocketTests.swift |  8 ++++----
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/Package.swift b/Package.swift
index b7c7ae2..bab9731 100644
--- a/Package.swift
+++ b/Package.swift
@@ -12,7 +12,7 @@ let package = Package(
         .library(name: "HummingbirdWSCompression", targets: ["HummingbirdWSCompression"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.0.0-beta.5"),
+        .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.0.0-beta.7"),
         .package(url: "https://github.com/apple/swift-http-types.git", from: "1.0.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.4.0"),
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.62.0"),
diff --git a/Sources/HummingbirdWebSocket/WebSocketRouter.swift b/Sources/HummingbirdWebSocket/WebSocketRouter.swift
index 5068895..928d1ed 100644
--- a/Sources/HummingbirdWebSocket/WebSocketRouter.swift
+++ b/Sources/HummingbirdWebSocket/WebSocketRouter.swift
@@ -59,17 +59,17 @@ public struct WebSocketHandlerReference<RequestContext: WebSocketRequestContext>
 }
 
 /// Request context protocol requirement for routers that support WebSockets
-public protocol WebSocketRequestContext: RequestContext, WebSocketContext {
+public protocol WebSocketRequestContext: InitializableFromSource<ApplicationRequestContextSource>, WebSocketContext {
     var webSocket: WebSocketHandlerReference<Self> { get }
 }
 
 /// Default implementation of a request context that supports WebSockets
-public struct BasicWebSocketRequestContext: WebSocketRequestContext {
-    public var coreContext: CoreRequestContext
+public struct BasicWebSocketRequestContext: RequestContext, WebSocketRequestContext {
+    public var coreContext: CoreRequestContextStorage
     public let webSocket: WebSocketHandlerReference<Self>
 
-    public init(channel: Channel, logger: Logger) {
-        self.coreContext = .init(allocator: channel.allocator, logger: logger)
+    public init(source: Source) {
+        self.coreContext = .init(source: source)
         self.webSocket = .init()
     }
 }
@@ -158,7 +158,7 @@ extension HTTP1WebSocketUpgradeChannel {
             let promise = channel.eventLoop.makePromise(of: ShouldUpgradeResult<WebSocketChannelHandler>.self)
             promise.completeWithTask {
                 let request = Request(head: head, body: .init(buffer: .init()))
-                let context = WSResponder.Context(channel: channel, logger: logger)
+                let context = WSResponder.Context(source: .init(channel: channel, logger: logger))
                 do {
                     let response = try await webSocketResponder.respond(to: request, context: context)
                     if response.status == .ok, let webSocketHandler = context.webSocket.handler.withLockedValue({ $0 }) {
diff --git a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
index 22f45bd..736fb96 100644
--- a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
+++ b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
@@ -477,13 +477,13 @@ final class HummingbirdWebSocketTests: XCTestCase {
 
     /// Test context from router is passed through to web socket
     func testRouterContextUpdate() async throws {
-        struct MyRequestContext: WebSocketRequestContext {
-            var coreContext: CoreRequestContext
+        struct MyRequestContext: RequestContext, WebSocketRequestContext {
+            var coreContext: CoreRequestContextStorage
             var webSocket: WebSocketHandlerReference<MyRequestContext>
             var name: String
 
-            init(channel: Channel, logger: Logger) {
-                self.coreContext = .init(allocator: channel.allocator, logger: logger)
+            init(source: Source) {
+                self.coreContext = .init(source: source)
                 self.webSocket = .init()
                 self.name = ""
             }