From 168cf5fbb9d90368d9e3c872214d946f7f0b3249 Mon Sep 17 00:00:00 2001
From: Paolo Rovelli <rovellipaolo@gmail.com>
Date: Wed, 20 Mar 2024 17:21:39 +0100
Subject: [PATCH] #45 Fix documentation preview with nested model

---
 README.md                         |  2 +-
 src/generateDocumentation.js      | 23 ++++++++++++++++++++++-
 src/generateDocumentation.spec.js |  6 ++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index d796fc8..ca989b3 100644
--- a/README.md
+++ b/README.md
@@ -316,7 +316,7 @@ To download it in YAML format, simply use `yml` or `yaml` extension in the "outp
 By default, the documentation will be generated in OpenAPI 2.0 (Swagger).
 To generate it in OpenAPI 3.0, use `oas30` or `openapi30` type in the "exportType" argument: `serverless generateDocumentation --outputFileName=filename.ext --exportType oas30`
 
-**NOTE:** The documentation preview generated locally might be invalid (e.g., when some mandatory fields are missing from the `serverless.yml` documentation) and might differt from the final documentation deployed to and downloaded from AWS.
+**NOTE:** The documentation preview generated locally might be invalid (e.g., when some mandatory fields are missing from the `serverless.yml` documentation) and might differ from the final documentation deployed to and downloaded from AWS.
 
 ### Download the documentation from AWS API Gateway
 
diff --git a/src/generateDocumentation.js b/src/generateDocumentation.js
index ee869cb..fe4c441 100644
--- a/src/generateDocumentation.js
+++ b/src/generateDocumentation.js
@@ -292,12 +292,33 @@ function generateModels(templateModels) {
 
   templateModels.forEach((templateModel) => {
     if (!templateModel.name || !templateModel.schema) return;
-    models[templateModel.name] = templateModel.schema;
+    models[templateModel.name] = generateModelSchema(templateModel.schema);
   });
 
   return models;
 }
 
+function generateModelSchema(templateSchema) {
+  if (!templateSchema) return templateSchema;
+
+  let schema = {};
+  Object.keys(templateSchema).forEach((key) => {
+    let field = templateSchema[key];
+    if (key === "$ref") {
+      const match = /{{\s*model\s*:\s*([\-\w]+)\s*}}/.exec(field);
+      if (match) {
+        field = `#/components/schemas/${match[1]}`;
+      }
+    } else if (!Array.isArray(field) && typeof field === "object" && field) {
+      // NOTE: We are only interested in looping through "items" and "properties" fields (which are both objects)!
+      field = generateModelSchema(field);
+    }
+    schema[key] = field;
+  });
+
+  return schema;
+}
+
 function generateSchema(field, excludeKeys) {
   if (field.schema) return field.schema;
 
diff --git a/src/generateDocumentation.spec.js b/src/generateDocumentation.spec.js
index 8aef2a5..70a51b3 100644
--- a/src/generateDocumentation.spec.js
+++ b/src/generateDocumentation.spec.js
@@ -346,6 +346,9 @@ describe('ServerlessAWSDocumentation', function () {
       ${[{ name: 'AnyModel', schema: { type: 'any-type' } }]}                                        | ${{ 'AnyModel': { type: 'any-type' } }}
       ${[{ name: 'AnyModel', contentType: 'application/json', schema: { type: 'any-type' } }]}       | ${{ 'AnyModel': { type: 'any-type' } }}
       ${[{ name: 'AnyModel', schema: { type: 'any-type' } }, { name: 'AnyOtherModel', schema: {} }]} | ${{ 'AnyModel': { type: 'any-type' }, 'AnyOtherModel': {} }}
+      ${[{ name: 'AnyModel', schema: { type: 'array', items: { "$ref": "{{model: AnyOtherModel}}" } } }]}                                          | ${{ 'AnyModel': { type: 'array', items: { "$ref": "#/components/schemas/AnyOtherModel" } } }}
+      ${[{ name: 'AnyModel', schema: { type: 'object', required: ['id'], properties: { id: { "$ref": "{{model: AnyOtherModel}}" } } } }]}          | ${{ 'AnyModel': { type: 'object', required: ['id'], properties: { id: { "$ref": "#/components/schemas/AnyOtherModel" } } } }}
+      ${[{ name: 'AnyModel', schema: { type: 'object', properties: { ids: { type: 'array', items: { "$ref": "{{model: AnyOtherModel}}" } } } } }]} | ${{ 'AnyModel': { type: 'object', properties: { ids: { type: 'array', items: { "$ref": "#/components/schemas/AnyOtherModel" } } } } }}
     `(
       'generates definitions field when exportType: swagger and models: $models',
       async ({
@@ -375,6 +378,9 @@ describe('ServerlessAWSDocumentation', function () {
       ${[{ name: 'AnyModel', schema: { type: 'any-type' } }]}                                        | ${{ schemas: { 'AnyModel': { type: 'any-type' } }, securitySchemes: {} }}
       ${[{ name: 'AnyModel', contentType: 'application/json', schema: { type: 'any-type' } }]}       | ${{ schemas: { 'AnyModel': { type: 'any-type' } }, securitySchemes: {} }}
       ${[{ name: 'AnyModel', schema: { type: 'any-type' } }, { name: 'AnyOtherModel', schema: {} }]} | ${{ schemas: { 'AnyModel': { type: 'any-type' }, 'AnyOtherModel': {} }, securitySchemes: {} }}
+      ${[{ name: 'AnyModel', schema: { type: 'array', items: { "$ref": "{{model: AnyOtherModel}}" } } }]}                                          | ${{ schemas: { 'AnyModel': { type: 'array', items: { "$ref": "#/components/schemas/AnyOtherModel" } } }, securitySchemes: {} }}
+      ${[{ name: 'AnyModel', schema: { type: 'object', required: ['id'], properties: { id: { "$ref": "{{model: AnyOtherModel}}" } } } }]}          | ${{ schemas: { 'AnyModel': { type: 'object', required: ['id'], properties: { id: { "$ref": "#/components/schemas/AnyOtherModel" } } } }, securitySchemes: {} }}
+      ${[{ name: 'AnyModel', schema: { type: 'object', properties: { ids: { type: 'array', items: { "$ref": "{{model: AnyOtherModel}}" } } } } }]} | ${{ schemas: { 'AnyModel': { type: 'object', properties: { ids: { type: 'array', items: { "$ref": "#/components/schemas/AnyOtherModel" } }} } }, securitySchemes: {} }}
     `(
       'generates components.schemas field when exportType: oas30 and models: $models',
       async ({