From 462be71028ae52602eb51344d84b1b3ddc18234b Mon Sep 17 00:00:00 2001
From: WoH <wolfgang.hobmaier@sensoring.de>
Date: Fri, 23 Dec 2022 00:25:33 +0100
Subject: [PATCH] feat: Supporting different content-type responses (#79)

---
 index.js           | 29 ++++++++++++---
 test/index.test.js | 88 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 4 deletions(-)

diff --git a/index.js b/index.js
index 549e302..7adaeff 100644
--- a/index.js
+++ b/index.js
@@ -43,16 +43,37 @@ function fastifyResponseValidation (fastify, opts, next) {
     const statusCodes = {}
     for (const statusCode in schema) {
       const responseSchema = schema[statusCode]
-      statusCodes[statusCode] = ajv.compile(
-        getSchemaAnyway(responseSchema)
-      )
+
+      if (responseSchema.content !== undefined) {
+        statusCodes[statusCode] = {}
+        for (const mediaName in responseSchema.content) {
+          statusCodes[statusCode][mediaName] = ajv.compile(
+            getSchemaAnyway(responseSchema.content[mediaName].schema)
+          )
+        }
+      } else {
+        statusCodes[statusCode] = ajv.compile(
+          getSchemaAnyway(responseSchema)
+        )
+      }
     }
 
     return preSerialization
 
     function preSerialization (req, reply, payload, next) {
-      const validate = statusCodes[reply.statusCode] || statusCodes[(reply.statusCode + '')[0] + 'xx']
+      let validate = statusCodes[reply.statusCode] || statusCodes[(reply.statusCode + '')[0] + 'xx']
+
       if (validate !== undefined) {
+        // Per media type validation
+        if (validate.constructor === Object) {
+          const mediaName = reply.getHeader('content-type').split(';')[0]
+          if (validate[mediaName] == null) {
+            next(new Error(`No schema defined for media type ${mediaName}`))
+            return
+          }
+          validate = validate[mediaName]
+        }
+
         const valid = validate(payload)
         if (!valid) {
           const err = new Error(schemaErrorsText(validate.errors))
diff --git a/test/index.test.js b/test/index.test.js
index 8dc88f2..dad4504 100644
--- a/test/index.test.js
+++ b/test/index.test.js
@@ -103,6 +103,94 @@ test('Should check only the assigned status code', async t => {
   t.strictSame(JSON.parse(response.payload), { answer: '42' })
 })
 
+test('Should check media types', async t => {
+  const fastify = Fastify()
+  await fastify.register(plugin)
+
+  fastify.route({
+    method: 'GET',
+    url: '/',
+    schema: {
+      response: {
+        '2xx': {
+          content: {
+            'application/geo+json': {
+              schema: {
+                type: 'object',
+                properties: {
+                  answer: { type: 'number' }
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    handler: async (req, reply) => {
+      reply.header('Content-Type', 'application/not+json')
+      return { answer: 42 }
+    }
+  })
+
+  const response = await fastify.inject({
+    method: 'GET',
+    url: '/'
+  })
+
+  t.equal(response.statusCode, 500)
+  t.strictSame(JSON.parse(response.payload), {
+    statusCode: 500,
+    error: 'Internal Server Error',
+    message: 'No schema defined for media type application/not+json'
+  })
+})
+
+test('Should support media types', async t => {
+  const fastify = Fastify()
+  await fastify.register(plugin)
+
+  fastify.route({
+    method: 'GET',
+    url: '/',
+    schema: {
+      response: {
+        '2xx': {
+          content: {
+            'application/a+json': {
+              schema: {
+                type: 'object',
+                properties: {
+                  answer: { type: 'boolean' }
+                }
+              }
+            },
+            'application/b+json': {
+              schema: {
+                type: 'object',
+                properties: {
+                  answer: { type: 'number' }
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    handler: async (req, reply) => {
+      reply.header('Content-Type', 'application/b+json')
+      return { answer: 42 }
+    }
+  })
+
+  const response = await fastify.inject({
+    method: 'GET',
+    url: '/'
+  })
+
+  t.equal(response.statusCode, 200)
+  t.strictSame(JSON.parse(response.payload), { answer: 42 })
+})
+
 test('Should check anyOf Schema', async t => {
   const fastify = Fastify()
   await fastify.register(plugin)