first shot at getting correct http status codes

even in when sub page builder is async

Author: rodja

nicegui/page.py

@@ -145,10 +145,7 @@ async def decorated(*dec_args, **dec_kwargs) -> Response:
                     dec_kwargs['client'] = client
                 try:
                     result = func(*dec_args, **dec_kwargs)
-                    # NOTE: after building the page, there might be sub pages that have 404 -- and initial requests should send 404 status request in such cases
-                    sub_pages_elements = [e for e in client.elements.values() if isinstance(e, SubPages)]
-                    if any(sub_pages.has_404 for sub_pages in sub_pages_elements):
-                        raise HTTPException(404, f'{client.sub_pages_router.current_path} not found')
+                    await self._await_sub_pages_and_raise_404(client)
                 except Exception as e:
                     return create_error_page(e, request)
             if helpers.is_coroutine_function(func):
@@ -174,6 +171,7 @@ async def wait_for_result() -> Optional[Response]:
                 else:
                     result = None
                     task.add_done_callback(check_for_late_return_value)
+                await self._await_sub_pages_and_raise_404(client)
             if isinstance(result, Response):  # NOTE if setup returns a response, we don't need to render the page
                 return result
             binding._refresh_step()  # pylint: disable=protected-access
@@ -192,3 +190,13 @@ async def wait_for_result() -> Optional[Response]:
         self.api_router.get(self._path, **self.kwargs)(decorated)
         Client.page_routes[func] = self.path
         return func
+
+    @staticmethod
+    async def _await_sub_pages_and_raise_404(client: Client) -> None:
+        """Some sub pages might finish async and decide 404"""
+        sub_pages_elements = [e for e in client.elements.values() if isinstance(e, SubPages)]
+        if any(sp._active_tasks for sp in sub_pages_elements):  # pylint: disable=protected-access
+            await asyncio.sleep(0)  # NOTE: give background tasks a brief chance to schedule nested sub pages
+        sub_pages_elements = [e for e in client.elements.values() if isinstance(e, SubPages)]
+        if any(sp.has_404 for sp in sub_pages_elements):
+            raise HTTPException(404, f'{client.sub_pages_router.current_path} not found')

tests/test_sub_pages.py

@@ -1087,8 +1087,12 @@ def main():
         ui.label('main page')
 
     screen.open('/bad_path')
-    screen.should_contain('HTTPException: 404: /bad_path not found')
     assert httpx.get(f'http://localhost:{Screen.PORT}/bad_path').status_code == 404
+    screen.should_contain('HTTPException: 404: /bad_path not found')
+
+    assert httpx.get(f'http://localhost:{Screen.PORT}/').status_code == 200
+    screen.open('/')
+    screen.should_contain('main page')
 
 
 def test_http_404_on_initial_request_with_async_page_builder(screen: Screen):
@@ -1103,9 +1107,38 @@ def main():
         ui.label('main page')
 
     screen.open('/bad_path')
+    assert httpx.get(f'http://localhost:{Screen.PORT}/bad_path').status_code == 404
     # NOTE: due to the async page builder, sub pages can not determine 404 status on initial request (see https://github.com/zauberzeug/nicegui/pull/5089)
-    screen.should_contain('404: sub page /bad_path not found')
-    assert httpx.get(f'http://localhost:{Screen.PORT}/bad_path').status_code == 200
+    screen.should_contain('HTTPException: 404: /bad_path not found')
+
+    assert httpx.get(f'http://localhost:{Screen.PORT}/').status_code == 200
+    screen.open('/')
+    screen.should_contain('main page')
+
+
+def test_http_404_on_initial_request_with_async_sub_page_builder(screen: Screen):
+    @ui.page('/')
+    @ui.page('/{_:path}')
+    def index():
+        ui.sub_pages({
+            '/': main,
+        })
+
+    async def main():
+        ui.label('main page')
+        ui.sub_pages({
+            '/': lambda: ui.label('sub main page'),
+            '/sub': lambda: ui.label('sub sub page'),
+        })
+
+    screen.open('/bad_path')
+    assert httpx.get(f'http://localhost:{Screen.PORT}/bad_path').status_code == 404
+    # NOTE: due to the async page builder, sub pages can not determine 404 status on initial request (see https://github.com/zauberzeug/nicegui/pull/5089)
+    screen.should_contain('HTTPException: 404: /bad_path not found')
+
+    assert httpx.get(f'http://localhost:{Screen.PORT}/sub').status_code == 200
+    screen.open('/sub')
+    screen.should_contain('sub sub page')
 
 
 def test_clearing_sub_pages_element(screen: Screen):