feat: add support for multiple headers

unjs · Jul 26, 2023 · 04fa824 · 04fa824
1 parent 57192c5
commit 04fa824
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/src/utils/response.ts b/src/utils/response.ts
@@ -148,9 +148,23 @@ export function appendResponseHeader(
 
 export const appendHeader = appendResponseHeader;
 
-export function removeResponseHeaders(event: H3Event): void {
-  for (const [name] of Object.entries(getResponseHeaders(event))) {
-    removeResponseHeader(event, name);
+/**
+ * Remove all response headers, or only those specified in the headerNames array.
+ * @param event H3 event
+ * @param headerNames Array of header names to remove
+ */
+export function removeResponseHeaders(
+  event: H3Event,
+  headerNames?: string[]
+): void {
+  if (headerNames && headerNames.length > 0) {
+    for (const name of headerNames) {
+      removeResponseHeader(event, name);
+    }
+  } else {
+    for (const [name] of Object.entries(getResponseHeaders(event))) {
+      removeResponseHeader(event, name);
+    }
   }
 }
 

diff --git a/test/header.test.ts b/test/header.test.ts
@@ -271,7 +271,7 @@ describe("", () => {
   });
 
   describe("removeResponseHeaders", () => {
-    it("can remove multiple response headers", async () => {
+    it("can remove all response headers", async () => {
       app.use(
         "/",
         eventHandler((event) => {
@@ -285,16 +285,34 @@ describe("", () => {
       expect(result.headers["header-2"]).toBeUndefined();
     });
 
-    it("can remove a single response header", async () => {
+    it("can remove multiple response headers", async () => {
       app.use(
         "/",
         eventHandler((event) => {
           appendHeader(event, "header-3", "3");
-          removeResponseHeader(event, "header-3");
+          appendHeader(event, "header-4", "4");
+          appendHeader(event, "header-5", "5");
+          removeResponseHeaders(event, ["header-3", "header-5"]);
         })
       );
       const result = await request.get("/");
       expect(result.headers["header-3"]).toBeUndefined();
+      expect(result.headers["header-4"]).toBe("4");
+      expect(result.headers["header-5"]).toBeUndefined();
+    });
+
+    it("can remove a single response header", async () => {
+      app.use(
+        "/",
+        eventHandler((event) => {
+          appendHeader(event, "header-6", "6");
+          appendHeader(event, "header-7", "7");
+          removeResponseHeader(event, "header-6");
+        })
+      );
+      const result = await request.get("/");
+      expect(result.headers["header-6"]).toBeUndefined();
+      expect(result.headers["header-7"]).toBe("7");
     });
   });
 });