From 1cc2604f6c0474c051bcc50038c81e8825cd9cba Mon Sep 17 00:00:00 2001
From: Kushal Chordiya <kushal.chordiya@trilogy.com>
Date: Thu, 15 Jun 2023 12:44:56 +0530
Subject: [PATCH] add security spec and tests

---
 chalice_spec/docs.py  |  4 +++
 tests/test_chalice.py | 83 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/chalice_spec/docs.py b/chalice_spec/docs.py
index 5ad5239..50d1c3e 100644
--- a/chalice_spec/docs.py
+++ b/chalice_spec/docs.py
@@ -37,6 +37,7 @@ def __init__(
         request: Optional[Type[BaseModel]] = None,
         response: Optional[Union[Response, Type[BaseModel]]] = None,
         responses: Optional[List[Response]] = None,
+        security: Optional[List[Dict[str, List[str]]]] = None,
     ):
         self.summary = summary
         self.description = description
@@ -45,6 +46,7 @@ def __init__(
 
         self.content_types = content_types
         self.request = request
+        self.security = security
 
         if response and responses:
             raise TypeError("You must only pass one of response or responses")
@@ -176,6 +178,8 @@ def _build_operation_from_operation(
             operation["tags"] = method.tags
         if method.parameters:
             operation["parameters"] = method.parameters
+        if method.security:
+            operation["security"] = method.security
 
         return operation
 
diff --git a/tests/test_chalice.py b/tests/test_chalice.py
index 6099a8d..8cea56a 100644
--- a/tests/test_chalice.py
+++ b/tests/test_chalice.py
@@ -553,3 +553,86 @@ def get_post():
             }
         },
     }
+
+
+# Test 10: test security
+def test_security():
+    app, spec = setup_test()
+    bearer_auth = {"type": "http", "scheme": "bearer"}
+    spec.components.security_scheme("BearerAuth", bearer_auth)
+
+    @app.route(
+        "/security",
+        methods=["POST"],
+        docs=Docs(
+            post=Op(
+                request=AnotherSchema,
+                response=Resp(
+                    code=201, description="Updated successfully!", model=TestSchema
+                ),
+                security=[{"BearerAuth": []}]
+            )
+        ),
+    )
+    def test():
+        pass
+
+    assert spec.to_dict() == {
+        "paths": {
+            "/security": {
+                "post": {
+                    "requestBody": {
+                        "content": {
+                            "application/json": {
+                                "schema": {"$ref": "#/components/schemas/AnotherSchema"}
+                            }
+                        }
+                    },
+                    "security": [{"BearerAuth": []}],
+                    "tags": ["/security"],
+                    "responses": {
+                        "201": {
+                            "description": "Updated successfully!",
+                            "content": {
+                                "application/json": {
+                                    "schema": {
+                                        "$ref": "#/components/schemas/TestSchema"
+                                    }
+                                }
+                            },
+                        }
+                    },
+                }
+            }
+        },
+        "info": {"title": "Test Schema", "version": "0.0.0"},
+        "openapi": "3.0.1",
+        "components": {
+            "schemas": {
+                "AnotherSchema": {
+                    "title": "AnotherSchema",
+                    "type": "object",
+                    "properties": {
+                        "nintendo": {"title": "Nintendo", "type": "string"},
+                        "atari": {"title": "Atari", "type": "string"},
+                    },
+                    "required": ["nintendo", "atari"],
+                },
+                "TestSchema": {
+                    "title": "TestSchema",
+                    "type": "object",
+                    "properties": {
+                        "hello": {"title": "Hello", "type": "string"},
+                        "world": {"title": "World", "type": "integer"},
+                    },
+                    "required": ["hello", "world"],
+                },
+            },
+            "securitySchemes": {
+                "BearerAuth": {
+                    "type": "http",
+                    "scheme": "bearer"
+                }
+            }
+        },
+    }