Skip to content

Commit

Permalink
🐛 Fix preserving route_class when calling include_router (#538)
Browse files Browse the repository at this point in the history
  • Loading branch information
dmontagu authored and tiangolo committed Oct 4, 2019
1 parent fdb6d43 commit dd96351
Show file tree
Hide file tree
Showing 2 changed files with 118 additions and 1 deletion.
5 changes: 4 additions & 1 deletion fastapi/routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -348,8 +348,10 @@ def add_api_route(
include_in_schema: bool = True,
response_class: Type[Response] = None,
name: str = None,
route_class_override: Optional[Type[APIRoute]] = None,
) -> None:
route = self.route_class(
route_class = route_class_override or self.route_class
route = route_class(
path,
endpoint=endpoint,
response_model=response_model,
Expand Down Expand Up @@ -487,6 +489,7 @@ def include_router(
include_in_schema=route.include_in_schema,
response_class=route.response_class or default_response_class,
name=route.name,
route_class_override=type(route),
)
elif isinstance(route, routing.Route):
self.add_route(
Expand Down
114 changes: 114 additions & 0 deletions tests/test_custom_route_class.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
import pytest
from fastapi import APIRouter, FastAPI
from fastapi.routing import APIRoute
from starlette.testclient import TestClient

app = FastAPI()


class APIRouteA(APIRoute):
x_type = "A"


class APIRouteB(APIRoute):
x_type = "B"


class APIRouteC(APIRoute):
x_type = "C"


router_a = APIRouter(route_class=APIRouteA)
router_b = APIRouter(route_class=APIRouteB)
router_c = APIRouter(route_class=APIRouteC)


@router_a.get("/")
def get_a():
return {"msg": "A"}


@router_b.get("/")
def get_b():
return {"msg": "B"}


@router_c.get("/")
def get_c():
return {"msg": "C"}


router_b.include_router(router=router_c, prefix="/c")
router_a.include_router(router=router_b, prefix="/b")
app.include_router(router=router_a, prefix="/a")


client = TestClient(app)

openapi_schema = {
"openapi": "3.0.2",
"info": {"title": "Fast API", "version": "0.1.0"},
"paths": {
"/a/": {
"get": {
"responses": {
"200": {
"description": "Successful Response",
"content": {"application/json": {"schema": {}}},
}
},
"summary": "Get A",
"operationId": "get_a_a__get",
}
},
"/a/b/": {
"get": {
"responses": {
"200": {
"description": "Successful Response",
"content": {"application/json": {"schema": {}}},
}
},
"summary": "Get B",
"operationId": "get_b_a_b__get",
}
},
"/a/b/c/": {
"get": {
"responses": {
"200": {
"description": "Successful Response",
"content": {"application/json": {"schema": {}}},
}
},
"summary": "Get C",
"operationId": "get_c_a_b_c__get",
}
},
},
}


@pytest.mark.parametrize(
"path,expected_status,expected_response",
[
("/a", 200, {"msg": "A"}),
("/a/b", 200, {"msg": "B"}),
("/a/b/c", 200, {"msg": "C"}),
("/openapi.json", 200, openapi_schema),
],
)
def test_get_path(path, expected_status, expected_response):
response = client.get(path)
assert response.status_code == expected_status
assert response.json() == expected_response


def test_route_classes():
routes = {}
r: APIRoute
for r in app.router.routes:
routes[r.path] = r
assert routes["/a/"].x_type == "A"
assert routes["/a/b/"].x_type == "B"
assert routes["/a/b/c/"].x_type == "C"

0 comments on commit dd96351

Please sign in to comment.