schema hadling changes, move tests to pytest

Author: tkfoss

pyproject.toml

@@ -27,6 +27,13 @@ dependencies = [
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[dependency-groups]
+dev = [
+    "pytest>=8.3.5",
+    "pytest-mock>=3.14.1",
+    "requests-mock>=1.12.1",
+]
+
 [project.scripts]
 openapi-gen = "serverless_openapi_generator.openapi_generator:main"
 openapi-validate = "openapi_spec_validator.__main__:main"

src/serverless_openapi_generator/openapi_generator.py

@@ -188,11 +188,30 @@ def create_operation_object(self, documentation, func, http_event):
             obj['requestBody'] = self.create_request_body(documentation.get('requestBody'))
         elif http_event.get('request', {}).get('schemas'):
             # Handle request schemas defined directly on the event
-            request_body_doc = {
-                'description': 'Request body inferred from event schema',
-                'requestModels': http_event['request']['schemas']
-            }
-            obj['requestBody'] = self.create_request_body(request_body_doc)
+            schemas = http_event['request']['schemas']
+            content = {}
+            for media_type, schema_info in schemas.items():
+                if isinstance(schema_info, str):
+                    schema_name = schema_info
+                    # Find the model name from the standardized models
+                    model = next((model for model in self.schema_handler.models if model.get('key') == schema_name), None)
+                    if model:
+                        content[media_type] = {
+                            'schema': {
+                                '$ref': f"#/components/schemas/{model['name']}"
+                            }
+                        }
+                elif isinstance(schema_info, dict):
+                    schema_name = ''.join(word.capitalize() for word in re.split(r'[/_-]', media_type))
+                    schema_ref = self.schema_handler.create_schema(schema_name, schema_info)
+                    content[media_type] = {'schema': {'$ref': schema_ref}}
+
+            if content:
+                obj['requestBody'] = {
+                    'description': 'Request body inferred from event schema',
+                    'content': content,
+                    'required': True
+                }
 
         # Handle private endpoints
         if http_event.get('private') is True:
@@ -264,24 +283,16 @@ def create_responses(self, documentation):
         return responses
 
     def create_media_type_object(self, models):
-        media_type_obj = {}
-        for media_type, model_info in models.items():
-            # Handle both simple model names and complex objects with schema
-            if isinstance(model_info, str):
-                model_name = model_info
-            elif isinstance(model_info, dict) and 'name' in model_info:
-                model_name = model_info['name']
-            else:
-                # Fallback for inline schemas, though we'll focus on named models
-                continue
-            
-            if model_name in self.open_api['components']['schemas']:
-                media_type_obj[media_type] = {
-                    'schema': {
-                        '$ref': f"#/components/schemas/{model_name}"
-                    }
-                }
-        return media_type_obj
+        content = {}
+        if models:
+            for media_type, schema_name in models.items():
+                print(f"Schema name: {schema_name}")
+                print(f"Models: {self.schema_handler.models}")
+                model_info = next((model for model in self.schema_handler.models if model['name'] == schema_name), None)
+                if model_info:
+                    schema_ref = self.schema_handler.create_schema(schema_name, model_info.get('schema'))
+                    content[media_type] = {'schema': {'$ref': schema_ref}}
+        return content
 
     def create_request_body(self, request_body_doc):
         content = self.create_media_type_object(request_body_doc.get('requestModels', {}))

src/serverless_openapi_generator/schema_handler.py

@@ -37,8 +37,7 @@ def _resolve_file_references(self, value):
 
     def _standardize_models(self):
         documentation = self.serverless_config.get('custom', {}).get('documentation', {})
-
-        standardized_models = []
+        
         model_sources = []
         if 'models' in documentation:
             models_config = self._resolve_file_references(documentation['models'])
@@ -51,6 +50,15 @@ def _standardize_models(self):
         if 'modelsList' in documentation:
             model_sources.extend(self._resolve_file_references(documentation['modelsList']))
 
+        api_gateway_models = self.serverless_config.get('provider', {}).get('apiGateway', {}).get('request', {}).get('schemas', {})
+        if api_gateway_models:
+            for model_key, model_definition in api_gateway_models.items():
+                if 'name' not in model_definition:
+                    model_definition['name'] = model_key
+                model_definition['key'] = model_key
+                model_sources.append(model_definition)
+
+        standardized_models = []
         for model in model_sources:
             if not isinstance(model, dict) or 'name' not in model:
                 continue
@@ -61,6 +69,8 @@ def _standardize_models(self):
                 'name': model.get('name'),
                 'description': model.get('description', ''),
             }
+            if 'key' in model:
+                std_model['key'] = model['key']
             if 'contentType' in model and 'schema' in model:
                 std_model['contentType'] = model['contentType']
                 std_model['schema'] = model['schema']
@@ -163,27 +173,23 @@ def _resolve_schema_references(self, schema):
 
         schema = self._clean_schema(schema)
 
-        registry = Registry()
-        if "definitions" in schema:
-            for name, sub_schema in schema["definitions"].items():
-                resource = Resource.from_contents(self._clean_schema(sub_schema), default_specification=DRAFT4)
-                registry = registry.with_resource(f"#/definitions/{name}", resource)
+        resource = Resource.from_contents(schema, default_specification=DRAFT4)
+        registry = Registry().with_resource("root", resource).crawl()
+        resolver = registry.resolver(base_uri="root")
 
-        main_resource = Resource.from_contents(schema, default_specification=DRAFT4)
-        registry = registry.with_resource("root", main_resource)
-        resolver = registry.resolver("root")
         dereferenced_schema = self._recursive_dereference(schema, resolver)
 
         if isinstance(dereferenced_schema, dict):
             dereferenced_schema.pop("definitions", None)
+
         return dereferenced_schema
 
     def _recursive_dereference(self, node, resolver):
         if isinstance(node, dict):
             if "$ref" in node:
                 try:
                     resolved = resolver.lookup(node["$ref"])
-                    return self._recursive_dereference(resolved.contents, resolver)
+                    return self._recursive_dereference(resolved.contents, resolved.resolver)
                 except Exception:
                     return node
             return {k: self._recursive_dereference(v, resolver) for k, v in node.items()}