Merge pull request #128 from asteasolutions/feature/#116-add-support-…

…for-response-headers-using-zod-objects added support for zod object for response headers
asteasolutions · May 2, 2023 · 38b74cb · 38b74cb
2 parents 55ceddb + 31b412d
commit 38b74cb
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 12 deletions.
diff --git a/spec/routes/index.spec.ts b/spec/routes/index.spec.ts
@@ -1,8 +1,7 @@
-import { z, ZodSchema } from 'zod';
-import { OperationObject, PathItemObject } from 'openapi3-ts';
+import { z } from 'zod';
 import { OpenAPIGenerator } from '../../src/openapi-generator';
-import { OpenAPIRegistry, RouteConfig } from '../../src/openapi-registry';
-import { createTestRoute, registerSchema, testDocConfig } from '../lib/helpers';
+import { OpenAPIRegistry } from '../../src/openapi-registry';
+import { createTestRoute, testDocConfig } from '../lib/helpers';
 
 const routeTests = ({
   registerFunction,
@@ -158,6 +157,50 @@ const routeTests = ({
 
       expect(responses['204']).toEqual({ description: 'Success' });
     });
+
+    it('can generate response headers', () => {
+      const registry = new OpenAPIRegistry();
+
+      registry[registerFunction]({
+        method: 'get',
+        path: '/',
+        responses: {
+          204: {
+            description: 'Success',
+            headers: z.object({
+              'Set-Cookie': z.string().openapi({
+                example: 'token=test',
+                description: 'Some string value',
+                param: {
+                  description: 'JWT session cookie',
+                },
+              }),
+            }),
+          },
+        },
+      });
+
+      const document = new OpenAPIGenerator(
+        registry.definitions,
+        '3.0.0'
+      ).generateDocument(testDocConfig);
+      const responses = document[rootDocPath]?.['/'].get.responses;
+
+      expect(responses['204']).toEqual({
+        description: 'Success',
+        headers: {
+          'Set-Cookie': {
+            schema: {
+              type: 'string',
+              example: 'token=test',
+              description: 'Some string value',
+            },
+            description: 'JWT session cookie',
+            required: true,
+          },
+        },
+      });
+    });
   });
 
   it('can generate paths with multiple examples', () => {

diff --git a/src/openapi-generator.ts b/src/openapi-generator.ts
@@ -10,6 +10,8 @@ import {
   ResponseObject,
   ContentObject,
   DiscriminatorObject,
+  HeadersObject,
+  BaseParameterObject,
 } from 'openapi3-ts';
 import type {
   AnyZodObject,
@@ -313,6 +315,24 @@ export class OpenAPIGenerator {
     ];
   }
 
+  private generateSimpleParameter(
+    zodSchema: ZodSchema<any>
+  ): BaseParameterObject {
+    const metadata = this.getMetadata(zodSchema);
+    const paramMetadata = metadata?.metadata?.param;
+
+    const required =
+      !this.isOptionalSchema(zodSchema) && !zodSchema.isNullable();
+
+    const schema = this.generateSimpleSchema(zodSchema);
+
+    return {
+      schema,
+      required,
+      ...(paramMetadata ? this.buildParameterMetadata(paramMetadata) : {}),
+    };
+  }
+
   private generateParameter(zodSchema: ZodSchema<any>): ParameterObject {
     const metadata = this.getMetadata(zodSchema);
 
@@ -332,17 +352,12 @@ export class OpenAPIGenerator {
       });
     }
 
-    const required =
-      !this.isOptionalSchema(zodSchema) && !zodSchema.isNullable();
-
-    const schema = this.generateSimpleSchema(zodSchema);
+    const baseParameter = this.generateSimpleParameter(zodSchema);
 
     return {
+      ...baseParameter,
       in: paramLocation,
       name: paramName,
-      schema,
-      required,
-      ...(paramMetadata ? this.buildParameterMetadata(paramMetadata) : {}),
     };
   }
 
@@ -536,18 +551,45 @@ export class OpenAPIGenerator {
 
   private getResponse({
     content,
+    headers,
     ...rest
   }: ResponseConfig): ResponseObject | ReferenceObject {
     const responseContent = content
       ? { content: this.getBodyContent(content) }
       : {};
 
+    if (!headers) {
+      return {
+        ...rest,
+        ...responseContent,
+      };
+    }
+
+    const responseHeaders = this.getResponseHeaders(headers);
+
     return {
       ...rest,
+      headers: responseHeaders,
       ...responseContent,
     };
   }
 
+  private getResponseHeaders(
+    headers: HeadersObject | AnyZodObject
+  ): HeadersObject {
+    if (!isZodType(headers, 'ZodObject')) {
+      return headers;
+    }
+
+    const schemaShape = headers._def.shape();
+
+    const responseHeaders = mapValues(schemaShape, _ =>
+      this.generateSimpleParameter(_)
+    );
+
+    return responseHeaders;
+  }
+
   private getBodyContent(content: ZodContentObject): ContentObject {
     return mapValues(content, config => {
       if (!isAnyZodType(config.schema)) {

diff --git a/src/openapi-registry.ts b/src/openapi-registry.ts
@@ -48,7 +48,7 @@ export interface ZodRequestBody {
 
 export interface ResponseConfig {
   description: string;
-  headers?: HeadersObject;
+  headers?: AnyZodObject | HeadersObject;
   links?: LinksObject;
   content?: ZodContentObject;
 }