Path Resolver For Route Handler Methods On A Controller (#480)

* #455 * removed memory reference * added temporary test Co-authored-by: Na'aman Hirschfeld <[email protected]>
litestar-org · Sep 15, 2022 · 3156a70 · 3156a70
1 parent 409296a
commit 3156a70
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 8 deletions.
diff --git a/starlite/handlers/base.py b/starlite/handlers/base.py
@@ -86,7 +86,11 @@ def __init__(
         self.opt: Dict[str, Any] = opt or {}
         self.owner: Optional[Union["Controller", "Router"]] = None
         self.signature_model: Optional[Type["SignatureModel"]] = None
-        self.paths = {normalize_path(p) for p in path} if path and isinstance(path, list) else {normalize_path(path or "/")}  # type: ignore
+        self.paths = (
+            {normalize_path(p) for p in path}
+            if path and isinstance(path, list)
+            else {normalize_path(path or "/")}  # type: ignore
+        )
 
     @property
     def dependency_name_set(self) -> Set[str]:

diff --git a/starlite/router.py b/starlite/router.py
@@ -1,3 +1,4 @@
+import collections
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -208,18 +209,18 @@ def route_handler_method_map(
     @property
     def route_handler_method_view(
         self,
-    ) -> Dict[Union[ASGIRouteHandler, "HttpMethod", HTTPRouteHandler, WebsocketRouteHandler], str]:
+    ) -> Dict[str, List[str]]:
         """
         Returns:
              A dictionary mapping route handlers to paths
         """
-        route_map: Dict[Union[ASGIRouteHandler, "HttpMethod", HTTPRouteHandler, WebsocketRouteHandler], str] = {}
+        route_map: Dict[str, List[str]] = collections.defaultdict(list)
         for route in self.routes:
             if isinstance(route, HTTPRoute):
                 for route_handler in route.route_handlers:
-                    route_map[route_handler] = route.path
+                    route_map[route_handler.fn.__qualname__].append(route.path)  # type: ignore
             else:
-                route_map[cast("WebSocketRoute", route).route_handler] = route.path
+                route_map[cast("WebSocketRoute", route).route_handler.fn.__qualname__].append(route.path)  # type: ignore
         return route_map
 
     @staticmethod

diff --git a/tests/test_router_registration.py b/tests/test_router_registration.py
@@ -5,10 +5,12 @@
     HttpMethod,
     ImproperlyConfiguredException,
     Router,
+    Starlite,
     WebSocket,
     get,
     patch,
     post,
+    put,
 )
 from starlite import route as route_decorator
 from starlite import websocket
@@ -125,9 +127,38 @@ def test_register_router_on_itself() -> None:
 
 
 def test_route_handler_method_view() -> None:
-    @get(path="/test")
+    @get(path="/root")
     def handler() -> None:
         ...
 
-    router = Router(path="/", route_handlers=[MyController, handler])
-    assert router.route_handler_method_view[handler] == "/test"
+    def _handler() -> None:
+        ...
+
+    put_handler = put("/modify")(_handler)
+    post_handler = post("/send")(_handler)
+
+    first_router = Router(path="/first", route_handlers=[MyController, handler, post_handler, put_handler])
+    second_router = Router(path="/second", route_handlers=[MyController, handler, post_handler, put_handler])
+
+    # Test routers.
+    assert first_router.route_handler_method_view[handler.fn.__qualname__] == ["/first/root"]  # type: ignore
+    assert second_router.route_handler_method_view[MyController.get_method.fn.__qualname__] == ["/second/test"]  # type: ignore
+
+    app = Starlite(route_handlers=[first_router, second_router])
+    # Test on Starlite instance.
+    assert app.route_handler_method_view[MyController.ws.fn.__qualname__] == [  # type: ignore
+        "/first/test/socket",
+        "/second/test/socket",
+    ]
+    assert app.route_handler_method_view[put_handler.fn.__qualname__] == [  # type: ignore
+        "/first/send",
+        "/first/modify",
+        "/second/send",
+        "/second/modify",
+    ]
+    assert app.route_handler_method_view[put_handler.fn.__qualname__] == [  # type: ignore
+        "/first/send",
+        "/first/modify",
+        "/second/send",
+        "/second/modify",
+    ]