localstack · calvernaz · Jul 2, 2023 · Jun 26, 2023
diff --git a/localstack/services/apigateway/helpers.py b/localstack/services/apigateway/helpers.py
@@ -22,6 +22,8 @@
 from localstack.aws.api.apigateway import (
     Authorizer,
     ConnectionType,
+    DocumentationPart,
+    DocumentationPartLocation,
     IntegrationType,
     Model,
     RequestValidator,
@@ -778,6 +780,25 @@ def apply_json_patch_safe(subject, patch_operations, in_place=True, return_list=
     return (results or [subject])[-1]
 
 
+def add_documentation_parts(rest_api_container, documentation):
+    for doc_part in documentation.get("documentationParts", []):
+        entity_id = short_uid()[:6]
+        location = doc_part["location"]
+        rest_api_container.documentation_parts[entity_id] = DocumentationPart(
+            id=entity_id,
+            location=DocumentationPartLocation(
+                type=location.get("type"),
+                path=location.get("path", "/")
+                if location.get("type") not in ["API", "MODEL"]
+                else None,
+                method=location.get("method"),
+                statusCode=location.get("statusCode"),
+                name=location.get("name"),
+            ),
+            properties=doc_part["properties"],
+        )
+
+
 def import_api_from_openapi_spec(
     rest_api: RestAPI, body: Dict, query_params: Dict, account_id: str = None, region: str = None
 ) -> Optional[RestAPI]:
@@ -1227,6 +1248,9 @@ def create_method_resource(child, method, method_schema):
     if api_key_source is not None:
         rest_api.api_key_source = api_key_source.upper()
 
+    documentation = resolved_schema.get(OpenAPIExt.DOCUMENTATION)
+    if documentation:
+        add_documentation_parts(rest_api_container, documentation)
     return rest_api
 
 

diff --git a/localstack/services/apigateway/provider.py b/localstack/services/apigateway/provider.py
@@ -31,6 +31,7 @@
     CreateAuthorizerRequest,
     CreateRestApiRequest,
     DocumentationPart,
+    DocumentationPartIds,
     DocumentationPartLocation,
     DocumentationParts,
     ExportResponse,
@@ -51,6 +52,7 @@
     NullableInteger,
     PutIntegrationRequest,
     PutIntegrationResponseRequest,
+    PutMode,
     PutRestApiRequest,
     RequestValidator,
     RequestValidators,
@@ -71,13 +73,15 @@
     EMPTY_MODEL,
     ERROR_MODEL,
     OpenApiExporter,
+    OpenAPIExt,
     apply_json_patch_safe,
     get_apigateway_store,
     import_api_from_openapi_spec,
     is_greedy_path,
     is_variable_path,
     log_template,
     multi_value_dict_for_list,
+    resolve_references,
 )
 from localstack.services.apigateway.invocations import invoke_rest_api_from_request
 from localstack.services.apigateway.models import RestApiContainer
@@ -945,6 +949,41 @@ def delete_documentation_part(
         if rest_api_container:
             rest_api_container.documentation_parts.pop(documentation_part_id, None)
 
+    def import_documentation_parts(
+        self,
+        context: RequestContext,
+        rest_api_id: String,
+        body: IO[Blob],
+        mode: PutMode = None,
+        fail_on_warnings: Boolean = None,
+    ) -> DocumentationPartIds:
+
+        body_data = body.read()
+        openapi_spec = parse_json_or_yaml(to_str(body_data))
+
+        store = get_apigateway_store(account_id=context.account_id, region=context.region)
+        if not (rest_api_container := store.rest_apis.get(rest_api_id)):
+            raise NotFoundException(
+                f"Invalid API identifier specified {context.account_id}:{rest_api_id}"
+            )
+
+        # https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-documenting-api-quick-start-import-export.html
+        resolved_schema = resolve_references(openapi_spec, rest_api_id=rest_api_id)
+        documentation = resolved_schema.get(OpenAPIExt.DOCUMENTATION)
+
+        ids = []
+        # overwrite mode
+        if mode == PutMode.overwrite:
+            rest_api_container.documentation_parts.clear()
+            for doc_part in documentation["documentationParts"]:
+                entity_id = short_uid()[:6]
+                rest_api_container.documentation_parts[entity_id] = DocumentationPart(
+                    id=entity_id, **doc_part
+                )
+                ids.append(entity_id)
+        # TODO: implement the merge mode
+        return DocumentationPartIds(ids=ids)
+
     # base path mappings
 
     def get_base_path_mappings(
@@ -1642,14 +1681,13 @@ def delete_model(
 
 def get_moto_rest_api(context: RequestContext, rest_api_id: str) -> MotoRestAPI:
     moto_backend = apigw_models.apigateway_backends[context.account_id][context.region]
-    rest_api = moto_backend.apis.get(rest_api_id)
-    if not rest_api:
+    if rest_api := moto_backend.apis.get(rest_api_id):
+        return rest_api
+    else:
         raise NotFoundException(
             f"Invalid API identifier specified {context.account_id}:{rest_api_id}"
         )
 
-    return rest_api
-
 
 def remove_empty_attributes_from_rest_api(rest_api: RestApi, remove_tags=True) -> RestApi:
     if not rest_api.get("binaryMediaTypes"):

diff --git a/tests/integration/apigateway/test_apigateway_api.py b/tests/integration/apigateway/test_apigateway_api.py
@@ -1,14 +1,17 @@
 import json
 import logging
+import os.path
 import time
 from operator import itemgetter
 
 import pytest
 from botocore.exceptions import ClientError
 
+from localstack.aws.api.apigateway import PutMode
 from localstack.services.apigateway.helpers import TAG_KEY_CUSTOM_ID
 from localstack.testing.aws.util import is_aws_cloud
 from localstack.testing.snapshots.transformer import KeyValueBasedTransformer, SortingTransformer
+from localstack.utils.files import load_file
 from localstack.utils.strings import short_uid
 from localstack.utils.sync import retry
 from tests.integration.apigateway.apigateway_fixtures import (
@@ -21,6 +24,9 @@
 
 LOG = logging.getLogger(__name__)
 
+PARENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+OAS_30_DOCUMENTATION_PARTS = os.path.join(PARENT_DIR, "files", "oas30_documentation_parts.json")
+
 
 @pytest.fixture(autouse=True)
 def apigw_snapshot_transformer(snapshot):
@@ -2049,3 +2055,36 @@ def test_invalid_delete_documentation_part(self, apigw_create_rest_api, snapshot
                 documentationPartId=documentation_part_id,
             )
         snapshot.match("delete_already_deleted_documentation_part", e.value.response)
+
+    @pytest.mark.aws_validated
+    def test_import_documentation_parts(self, aws_client, import_apigw, snapshot):
+        # snapshot array "ids"
+        snapshot.add_transformer(snapshot.transform.jsonpath("$..ids[*]", "id"))
+        # create api with documentation imports
+        spec_file = load_file(OAS_30_DOCUMENTATION_PARTS)
+        response, root_id = import_apigw(body=spec_file, failOnWarnings=True)
+        rest_api_id = response["id"]
+
+        # get documentation parts to make sure import worked
+        response = aws_client.apigateway.get_documentation_parts(restApiId=rest_api_id)
+        snapshot.match("create-import-documentations_parts", response["items"])
+
+        # delete documentation parts
+        for doc_part_item in response["items"]:
+            response = aws_client.apigateway.delete_documentation_part(
+                restApiId=rest_api_id,
+                documentationPartId=doc_part_item["id"],
+            )
+            assert response["ResponseMetadata"]["HTTPStatusCode"] == 202
+
+        # make sure delete parts are gone
+        response = aws_client.apigateway.get_documentation_parts(restApiId=rest_api_id)
+        assert len(response["items"]) == 0
+
+        # import documentation parts using import documentation parts api
+        response = aws_client.apigateway.import_documentation_parts(
+            restApiId=rest_api_id,
+            mode=PutMode.overwrite,
+            body=spec_file,
+        )
+        snapshot.match("import-documentation-parts", response)
diff --git a/tests/integration/apigateway/test_apigateway_api.snapshot.json b/tests/integration/apigateway/test_apigateway_api.snapshot.json
@@ -2538,5 +2538,70 @@
         }
       }
     }
+  },
+  "tests/integration/apigateway/test_apigateway_api.py::TestApiGatewayApiDocumentationPart::test_import_documentation_parts": {
+    "recorded-date": "26-06-2023, 12:01:38",
+    "recorded-content": {
+      "create-import-documentations_parts": [
+        {
+          "id": "<id:1>",
+          "location": {
+            "type": "API"
+          },
+          "properties": {
+            "description": "API description",
+            "info": {
+              "description": "API info description 4",
+              "version": "API info version 3"
+            }
+          }
+        },
+        {
+          "id": "<id:2>",
+          "location": {
+            "type": "METHOD",
+            "path": "/",
+            "method": "GET"
+          },
+          "properties": {
+            "description": "Method description."
+          }
+        },
+        {
+          "id": "<id:3>",
+          "location": {
+            "type": "MODEL",
+            "name": "<name:1>"
+          },
+          "properties": {
+            "title": "<name:1> Schema"
+          }
+        },
+        {
+          "id": "<id:4>",
+          "location": {
+            "type": "RESPONSE",
+            "path": "/",
+            "method": "GET",
+            "statusCode": "200"
+          },
+          "properties": {
+            "description": "200 response"
+          }
+        }
+      ],
+      "import-documentation-parts": {
+        "ids": [
+          "<id:5>",
+          "<id:6>",
+          "<id:7>",
+          "<id:8>"
+        ],
+        "ResponseMetadata": {
+          "HTTPHeaders": {},
+          "HTTPStatusCode": 200
+        }
+      }
+    }
   }
 }
diff --git a/tests/integration/files/oas30_documentation_parts.json b/tests/integration/files/oas30_documentation_parts.json
@@ -0,0 +1,85 @@
+{
+  "openapi": "3.0.0",
+  "info": {
+    "description": "description",
+    "version": "1",
+    "title": "doc"
+  },
+  "paths": {
+    "/": {
+      "get": {
+        "description": "Method description.",
+        "responses": {
+          "200": {
+            "description": "200 response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/Empty"
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  },
+  "x-amazon-apigateway-documentation": {
+    "version": "1.0.3",
+    "documentationParts": [
+      {
+        "location": {
+          "type": "API"
+        },
+        "properties": {
+          "description": "API description",
+          "info": {
+            "description": "API info description 4",
+            "version": "API info version 3"
+          }
+        }
+      },
+      {
+        "location": {
+          "type": "METHOD",
+          "method": "GET"
+        },
+        "properties": {
+          "description": "Method description."
+        }
+      },
+      {
+        "location": {
+          "type": "MODEL",
+          "name": "Empty"
+        },
+        "properties": {
+          "title": "Empty Schema"
+        }
+      },
+      {
+        "location": {
+          "type": "RESPONSE",
+          "method": "GET",
+          "statusCode": "200"
+        },
+        "properties": {
+          "description": "200 response"
+        }
+      }
+    ]
+  },
+  "servers": [
+    {
+      "url": "/"
+    }
+  ],
+  "components": {
+    "schemas": {
+      "Empty": {
+        "type": "object",
+        "title": "Empty Schema"
+      }
+    }
+  }
+}