Express + Zod to OpenAPI

openapi.test.ts

const mockRegistry = { definitions: {}, registerPath: jest.fn() };
const mockGenerator = { generateDocument: jest.fn() };

jest.mock('@asteasolutions/zod-to-openapi', () => ({
  extendZodWithOpenApi: jest.fn(),
  OpenAPIRegistry: function () {
    return mockRegistry;
  },
  OpenApiGeneratorV3: function () {
    return mockGenerator;
  },
}));
jest.mock('node:fs/promises');
jest.mock('swagger-ui-express');
jest.mock('yaml');

import { Response } from 'express';
import fs from 'node:fs/promises';
import swaggerUi from 'swagger-ui-express';
import yaml from 'yaml';
import { z } from 'zod';

import * as openapi from '@lib/openapi';
import { mocked } from '@lib/test-utils';

describe('openapi', () => {
  describe('publishDocs', () => {
    const mockApp = { use: jest.fn(), get: jest.fn() };
    const mockConfig = {
      openapi: '3.0.0',
      info: {
        title: 'Test API',
        version: '1.0.0',
      },
    };
    const filePath = '/path/to/swagger.yml';

    it('should generate and publish Swagger documentation', async () => {
      const mockDocument = { ...mockConfig, paths: { '/test': {} } };
      const mockHandler = jest.fn();
      const mockResponse = { json: jest.fn() } as unknown as Response;

      mockGenerator.generateDocument.mockReturnValueOnce(mockDocument);
      mocked(yaml.stringify).mockReturnValueOnce('mocked_yaml_content');
      mocked(swaggerUi.setup).mockReturnValueOnce(mockHandler);

      await openapi.publishDocs(mockConfig, filePath, openapi.registry, mockApp as never);

      // Call the registered handler
      mockApp.get.mock.calls[0][1]({}, mockResponse, jest.fn());

      expect(yaml.stringify).toHaveBeenCalledWith(mockDocument);
      expect(fs.writeFile).toHaveBeenCalledWith(filePath, 'mocked_yaml_content', { encoding: 'utf-8' });
      expect(mockApp.use).toHaveBeenCalledWith(`/api-docs`, [], expect.any(Function));
      expect(mockApp.get).toHaveBeenCalledWith('/swagger.json', expect.any(Function));
      expect(mockResponse.json).toHaveBeenCalledWith(mockDocument);
    });

    it('should handle errors during document generation and writing', async () => {
      const mockError = new Error('Test error');

      mocked(yaml.stringify).mockImplementationOnce(() => {
        throw mockError;
      });

      await expect(openapi.publishDocs(mockConfig, filePath, openapi.registry, mockApp as never)).rejects.toThrow(
        mockError,
      );

      expect(fs.writeFile).not.toHaveBeenCalled();
      expect(mockApp.use).not.toHaveBeenCalled();
      expect(mockApp.get).not.toHaveBeenCalled();
    });
  });

  describe('registerRoute', () => {
    it('should register the route and handle requests without input', () => {
      const router = { get: jest.fn() };
      const config = {
        method: 'get',
        path: '/test',
        responses: {
          200: {
            content: {
              'application/json': { schema: z.object({ message: z.string() }) },
            },
          },
        },
      };
      const mockHandler = jest.fn();
      const mockNext = jest.fn();

      openapi.registerRoute(router as never, config as never, mockHandler);

      // Simulate a request without input
      router.get.mock.calls[0][1]({}, {}, mockNext);

      expect(mockHandler).toHaveBeenCalled();
      expect(mockNext).not.toHaveBeenCalled();
    });

    it('should register the route and handle requests with valid input', () => {
      const router = { get: jest.fn() };
      const config = {
        method: 'get',
        path: '/test',
        request: {
          body: z.object({ name: z.string() }),
          params: z.object({ id: z.string() }),
          query: z.object({ page: z.string() }),
        },
        responses: {
          200: {
            content: {
              'application/json': { schema: z.object({ message: z.string() }) },
            },
          },
        },
      };
      const mockHandler = jest.fn();
      const mockRequest = {
        body: { name: 'John' },
        params: { id: '123' },
        query: { page: '1' },
      } as unknown as Request;
      const mockNext = jest.fn();

      openapi.registerRoute(router as never, config as never, mockHandler);

      // Simulate a valid request
      router.get.mock.calls[0][1](mockRequest, {}, mockNext);

      expect(mockHandler).toHaveBeenCalled();
      expect(mockNext).not.toHaveBeenCalled();
    });

    it('should handle requests with invalid input', () => {
      const router = { get: jest.fn() };
      const config = {
        method: 'get',
        path: '/test',
        request: {
          params: z.object({ id: z.string() }),
        },
        responses: {
          200: {
            content: {
              'application/json': { schema: z.object({ message: z.string() }) },
            },
          },
        },
      };
      const mockHandler = jest.fn();
      const mockNext = jest.fn();

      openapi.registerRoute(router as never, config as never, mockHandler);

      // Simulate an invalid request (missing required parameter)
      router.get.mock.calls[0][1]({}, {}, mockNext);

      expect(mockHandler).not.toHaveBeenCalled();
      expect(mockNext).toHaveBeenCalledWith(expect.any(Error));
    });
  });
});

openapi.ts

import {
  OpenAPIRegistry,
  OpenApiGeneratorV3,
  RouteConfig,
  ZodRequestBody,
  extendZodWithOpenApi,
} from '@asteasolutions/zod-to-openapi';
import { OpenAPIObjectConfig } from '@asteasolutions/zod-to-openapi/dist/v3.0/openapi-generator';
import { Application, NextFunction, Request, RequestHandler, Response, Router } from 'express';
import fs from 'node:fs/promises';
import swaggerUi from 'swagger-ui-express';
import yaml from 'yaml';
import { z } from 'zod';

extendZodWithOpenApi(z);

export const registry = new OpenAPIRegistry();

export async function publishDocs(
  config: OpenAPIObjectConfig,
  filePath: string,
  reg: OpenAPIRegistry,
  app: Application,
) {
  const generator = new OpenApiGeneratorV3(reg.definitions);
  const document = generator.generateDocument(config);
  const fileContent = yaml.stringify(document);

  await fs.writeFile(filePath, fileContent, { encoding: 'utf-8' });

  app.use(`/api-docs`, swaggerUi.serve, swaggerUi.setup(document));
  app.get('/swagger.json', (_req, res) => {
    res.json(document);
  });
}

type ZodObject<T extends z.ZodRawShape = z.ZodRawShape> = z.ZodObject<T>;

type RequestParams<T extends RouteConfig> = T['request'] extends { params: ZodObject }
  ? z.infer<T['request']['params']>
  : never;

type ResponseBody<T extends RouteConfig> = T['responses'][keyof T['responses']] extends {
  content: { [MediaType in keyof T['responses'][keyof T['responses']]['content']]: { schema: ZodObject } };
}
  ? {
      [MediaType in keyof T['responses'][keyof T['responses']]['content']]: z.infer<
        T['responses'][keyof T['responses']]['content'][MediaType]['schema']
      >;
    }[keyof T['responses'][keyof T['responses']]['content']]
  : never;

type RequestBody<T extends RouteConfig> = T['request'] extends { body: ZodRequestBody }
  ? T['request']['body'] extends {
      content: { [MediaType in keyof T['request']['body']['content']]: { schema: ZodObject } };
    }
    ? {
        [MediaType in keyof T['request']['body']['content']]: z.infer<
          T['request']['body']['content'][MediaType]['schema']
        >;
      }[keyof T['responses'][keyof T['responses']]['content']]
    : never
  : never;

type RequestQuery<T extends RouteConfig> = T['request'] extends { query: ZodObject }
  ? z.infer<T['request']['query']>
  : never;

export type Handler<T extends RouteConfig> = RequestHandler<
  RequestParams<T>,
  ResponseBody<T>,
  RequestBody<T>,
  RequestQuery<T>
>;

export function registerRoute<T extends RouteConfig>(router: Router, config: T, ...handlers: Handler<T>[]): void {
  registry.registerPath(config);

  const validatedHandlers = handlers.map((handler) => (req: Request, res: Response, next: NextFunction) => {
    try {
      const validatedRequest = validateRequest(config, req as never);
      handler(validatedRequest, res, next);
    } catch (err) {
      // TODO: log error
      next(new Error('Input validation failed'));
    }
  });

  router[config.method](config.path, ...validatedHandlers);
}

function validateRequest<
  T extends RouteConfig,
  R extends Request<RequestParams<T>, ResponseBody<T>, RequestBody<T>, RequestQuery<T>>,
>(config: T, request: Request): R {
  const result = { ...request };

  if (config.request?.params) {
    result.params = config.request.params.parse(request.params);
  }
  if (config.request?.body) {
    result.body = (config.request.body as unknown as ZodObject).parse(request.body);
  }
  if (config.request?.query) {
    result.query = config.request.query.parse(request.query);
  }

  return result as R;
}