From 178d4c44773dd579c59836cea6c6e0f94da51a02 Mon Sep 17 00:00:00 2001
From: Timo <38291523+lovetodream@users.noreply.github.com>
Date: Tue, 3 Sep 2024 08:01:01 +0200
Subject: [PATCH] add additionalHeaders to initializer + docc and tests

---
 .../HummingbirdElementary/HTMLResponse.swift   | 18 +++++++++++++++++-
 .../HTMLResponseTests.swift                    | 16 +++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/Sources/HummingbirdElementary/HTMLResponse.swift b/Sources/HummingbirdElementary/HTMLResponse.swift
index c191ed3..39ae460 100644
--- a/Sources/HummingbirdElementary/HTMLResponse.swift
+++ b/Sources/HummingbirdElementary/HTMLResponse.swift
@@ -26,15 +26,31 @@ public struct HTMLResponse<Content: HTML & Sendable>: Sendable {
     public var chunkSize: Int
 
     /// Response headers
+    ///
+    /// It can be used to add additional headers to a predefined set of fields.
+    ///
+    /// - Note: If a new set of headers is assigned, all predefined headers are removed.
+    ///
+    /// ```swift
+    /// var response = HTMLResponse { ... }
+    /// response.headers[.init("foo")!] = "bar"
+    /// return response
+    /// ```
     public var headers: HTTPFields = [.contentType: "text/html; charset=utf-8"]
 
     /// Creates a new HTMLResponse
     ///
     /// - Parameters:
     ///   - chunkSize: The number of bytes to write to the response body at a time.
+    ///   - additionalHeaders: Additional headers to be merged with predefined headers.
     ///   - content: The `HTML` content to render in the response.
-    public init(chunkSize: Int = 1024, @HTMLBuilder content: () -> Content) {
+    public init(chunkSize: Int = 1024, additionalHeaders: HTTPFields = [:], @HTMLBuilder content: () -> Content) {
         self.chunkSize = chunkSize
+        if additionalHeaders.contains(.contentType) {
+            self.headers = additionalHeaders
+        } else {
+            self.headers = [.contentType: "text/html; charset=utf-8"] + additionalHeaders
+        }
         self.content = content()
     }
 }
diff --git a/Tests/HummingbirdElementaryTests/HTMLResponseTests.swift b/Tests/HummingbirdElementaryTests/HTMLResponseTests.swift
index 3765864..64ac1b6 100644
--- a/Tests/HummingbirdElementaryTests/HTMLResponseTests.swift
+++ b/Tests/HummingbirdElementaryTests/HTMLResponseTests.swift
@@ -53,7 +53,7 @@ final class HTMLResponseTests: XCTestCase {
 
     func testRespondsWithCustomHeaders() async throws {
         let router = Router().get { _, _ in
-            var response = HTMLResponse { TestPage() }
+            var response = HTMLResponse(additionalHeaders: [.init("foo")!: "bar"]) { EmptyHTML() }
             response.headers[.init("hx-refresh")!] = "true"
             return response
         }
@@ -61,7 +61,21 @@ final class HTMLResponseTests: XCTestCase {
         try await Application(router: router).test(.router) { client in
             let response = try await client.execute(uri: "/", method: .get)
 
+            XCTAssertEqual(response.headers[.init("foo")!], "bar")
             XCTAssertEqual(response.headers[.init("hx-refresh")!], "true")
+            XCTAssertEqual(response.headers[.contentType], "text/html; charset=utf-8")
+        }
+    }
+
+    func testRespondsWithOverwrittenContentType() async throws {
+        let router = Router().get { _, _ in
+            HTMLResponse(additionalHeaders: [.contentType: "new"]) { EmptyHTML() }
+        }
+
+        try await Application(router: router).test(.router) { client in
+            let response = try await client.execute(uri: "/", method: .get)
+
+            XCTAssertEqual(response.headers[.contentType], "new")
         }
     }
 }