From 1dd98d2752b47a3e7b0dd48f5b4223af4f759135 Mon Sep 17 00:00:00 2001
From: Georg K <jar3b@outlook.com>
Date: Thu, 4 Aug 2022 23:01:04 +0300
Subject: [PATCH] feat: add support for definition in query param

---
 aiohttp_pydantic/oas/view.py          |  4 ++++
 tests/test_oas/test_view.py           | 33 ++++++++++++++++++++++-----
 tests/test_validation_query_string.py | 29 +++++++++++++++++++++++
 3 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/aiohttp_pydantic/oas/view.py b/aiohttp_pydantic/oas/view.py
index 37bd3ff..858d7c2 100644
--- a/aiohttp_pydantic/oas/view.py
+++ b/aiohttp_pydantic/oas/view.py
@@ -126,6 +126,10 @@ def _add_http_method_to_oas(
                 ref_template="#/components/schemas/{model}"
             )
 
+            # move definitions
+            if def_sub_schemas := oas_operation.parameters[i].schema.pop("definitions", None):
+                oas.components.schemas.update(def_sub_schemas)
+
     return_type = get_type_hints(handler).get("return")
     if return_type is not None:
         _OASResponseBuilder(oas, oas_operation, status_code_descriptions).build(
diff --git a/tests/test_oas/test_view.py b/tests/test_oas/test_view.py
index 040b579..e8f92c3 100644
--- a/tests/test_oas/test_view.py
+++ b/tests/test_oas/test_view.py
@@ -20,6 +20,11 @@ class Color(str, Enum):
     PINK = "pink"
 
 
+class Lang(str, Enum):
+    EN = 'en'
+    FR = 'fr'
+
+
 class Toy(BaseModel):
     name: str
     color: Color
@@ -33,7 +38,7 @@ class Pet(BaseModel):
 
 class PetCollectionView(PydanticView):
     async def get(
-        self, format: str, name: Optional[str] = None, *, promo: Optional[UUID] = None
+            self, format: str, lang: Lang = Lang.EN, name: Optional[str] = None, *, promo: Optional[UUID] = None
     ) -> r200[List[Pet]]:
         """
         Get a list of pets
@@ -51,11 +56,11 @@ class PetCollectionView(PydanticView):
 
 class PetItemView(PydanticView):
     async def get(
-        self,
-        id: int,
-        /,
-        size: Union[int, Literal["x", "l", "s"]],
-        day: Union[int, Literal["now"]] = "now",
+            self,
+            id: int,
+            /,
+            size: Union[int, Literal["x", "l", "s"]],
+            day: Union[int, Literal["now"]] = "now",
     ) -> Union[r200[Pet], r404]:
         return web.json_response()
 
@@ -116,6 +121,12 @@ async def test_generated_oas_should_have_components_schemas(generated_oas):
             "title": "Color",
             "type": "string",
         },
+        'Lang': {
+            'description': 'An enumeration.',
+            'enum': ['en', 'fr'],
+            'title': 'Lang',
+            'type': 'string'
+        },
         "Toy": {
             "properties": {
                 "color": {"$ref": "#/components/schemas/Color"},
@@ -143,6 +154,16 @@ async def test_pets_route_should_have_get_method(generated_oas):
                 "required": True,
                 "schema": {"title": "format", "type": "string"},
             },
+            {
+                'in': 'query',
+                'name': 'lang',
+                'required': False,
+                'schema': {
+                    'allOf': [{'$ref': '#/components/schemas/Lang'}],
+                    'default': 'en',
+                    'title': 'lang'
+                }
+            },
             {
                 "in": "query",
                 "name": "name",
diff --git a/tests/test_validation_query_string.py b/tests/test_validation_query_string.py
index a1b610c..411410a 100644
--- a/tests/test_validation_query_string.py
+++ b/tests/test_validation_query_string.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import Optional, List
 from pydantic import Field
 from aiohttp import web
@@ -54,6 +55,20 @@ class ArticleViewWithPaginationGroup(PydanticView):
         )
 
 
+class Lang(str, Enum):
+    EN = 'en'
+    FR = 'fr'
+
+
+class ArticleViewWithEnumInQuery(PydanticView):
+    async def get(self, lang: Lang):
+        return web.json_response(
+            {
+                "lang": lang
+            }
+        )
+
+
 async def test_get_article_without_required_qs_should_return_an_error_message(
     aiohttp_client, event_loop
 ):
@@ -236,6 +251,20 @@ async def test_get_article_with_page_and_page_size(aiohttp_client, event_loop):
     assert resp.content_type == "application/json"
 
 
+async def test_get_article_with_enum_in_query(aiohttp_client, event_loop):
+    app = web.Application()
+    app.router.add_view("/article", ArticleViewWithEnumInQuery)
+
+    client = await aiohttp_client(app)
+
+    resp = await client.get(
+        "/article", params={"lang": Lang.EN.value}
+    )
+    assert await resp.json() == {'lang': Lang.EN}
+    assert resp.status == 200
+    assert resp.content_type == "application/json"
+
+
 async def test_get_article_with_page_and_wrong_page_size(aiohttp_client, event_loop):
     app = web.Application()
     app.router.add_view("/article", ArticleViewWithPaginationGroup)