From 73dbbeab554b8d94fffd934d951009f026491015 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <divums@users.noreply.github.com>
Date: Fri, 30 Aug 2019 19:17:42 +0300
Subject: [PATCH] :sparkles: Allow additional responses to use status ranges
 and "default" (#435)

---
 fastapi/openapi/utils.py                  | 15 +++++++--
 tests/test_additional_responses_bad.py    | 40 +++++++++++++++++++++++
 tests/test_additional_responses_router.py | 20 ++++++++++--
 3 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_additional_responses_bad.py

diff --git a/fastapi/openapi/utils.py b/fastapi/openapi/utils.py
index 718c7de27..2ee8c241c 100644
--- a/fastapi/openapi/utils.py
+++ b/fastapi/openapi/utils.py
@@ -43,6 +43,15 @@ validation_error_response_definition = {
     },
 }
 
+status_code_ranges: Dict[str, str] = {
+    "1XX": "Information",
+    "2XX": "Success",
+    "3XX": "Redirection",
+    "4XX": "Client Error",
+    "5XX": "Server Error",
+    "default": "Default Response",
+}
+
 
 def get_openapi_params(dependant: Dependant) -> List[Field]:
     flat_dependant = get_flat_dependant(dependant, skip_repeats=True)
@@ -190,12 +199,14 @@ def get_openapi_path(
                         response.setdefault("content", {}).setdefault(
                             "application/json", {}
                         )["schema"] = response_schema
-                    status_text = http.client.responses.get(int(additional_status_code))
+                    status_text: Optional[str] = status_code_ranges.get(
+                        str(additional_status_code).upper()
+                    ) or http.client.responses.get(int(additional_status_code))
                     response.setdefault(
                         "description", status_text or "Additional Response"
                     )
                     operation.setdefault("responses", {})[
-                        str(additional_status_code)
+                        str(additional_status_code).upper()
                     ] = response
             status_code = str(route.status_code)
             response_schema = {"type": "string"}
diff --git a/tests/test_additional_responses_bad.py b/tests/test_additional_responses_bad.py
new file mode 100644
index 000000000..fda475576
--- /dev/null
+++ b/tests/test_additional_responses_bad.py
@@ -0,0 +1,40 @@
+import pytest
+from fastapi import FastAPI
+from starlette.testclient import TestClient
+
+app = FastAPI()
+
+
+@app.get("/a", responses={"hello": {"description": "Not a valid additional response"}})
+async def a():
+    pass  # pragma: no cover
+
+
+openapi_schema = {
+    "openapi": "3.0.2",
+    "info": {"title": "Fast API", "version": "0.1.0"},
+    "paths": {
+        "/a": {
+            "get": {
+                "responses": {
+                    # this is how one would imagine the openapi schema to be
+                    # but since the key is not valid, openapi.utils.get_openapi will raise ValueError
+                    "hello": {"description": "Not a valid additional response"},
+                    "200": {
+                        "description": "Successful Response",
+                        "content": {"application/json": {"schema": {}}},
+                    },
+                },
+                "summary": "A",
+                "operationId": "a_a_get",
+            }
+        }
+    },
+}
+
+client = TestClient(app)
+
+
+def test_openapi_schema():
+    with pytest.raises(ValueError):
+        client.get("/openapi.json")
diff --git a/tests/test_additional_responses_router.py b/tests/test_additional_responses_router.py
index 49ef5f049..ce66ead7e 100644
--- a/tests/test_additional_responses_router.py
+++ b/tests/test_additional_responses_router.py
@@ -10,12 +10,24 @@ async def a():
     return "a"
 
 
-@router.get("/b", responses={502: {"description": "Error 2"}})
+@router.get(
+    "/b",
+    responses={
+        502: {"description": "Error 2"},
+        "4XX": {"description": "Error with range, upper"},
+    },
+)
 async def b():
     return "b"
 
 
-@router.get("/c", responses={501: {"description": "Error 3"}})
+@router.get(
+    "/c",
+    responses={
+        "400": {"description": "Error with str"},
+        "5xx": {"description": "Error with range, lower"},
+    },
+)
 async def c():
     return "c"
 
@@ -43,6 +55,7 @@ openapi_schema = {
             "get": {
                 "responses": {
                     "502": {"description": "Error 2"},
+                    "4XX": {"description": "Error with range, upper"},
                     "200": {
                         "description": "Successful Response",
                         "content": {"application/json": {"schema": {}}},
@@ -55,7 +68,8 @@ openapi_schema = {
         "/c": {
             "get": {
                 "responses": {
-                    "501": {"description": "Error 3"},
+                    "400": {"description": "Error with str"},
+                    "5XX": {"description": "Error with range, lower"},
                     "200": {
                         "description": "Successful Response",
                         "content": {"application/json": {"schema": {}}},