Provide tests for run.io_bound and run.cpu_bound (#2234)

* provide tests for run.io_bound and run.cpu_bound

* fix import

* change test fixture from screen to fixture

* add test to ensure handling of unpickable exceptions in run.cpu_bound

* ensure we do not break the process pool with problematic exceptions

* code review

---------

Co-authored-by: Falko Schindler <falko@zauberzeug.com>

nicegui/run.py

@@ -1,5 +1,6 @@
 import asyncio
 import sys
+import traceback
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from functools import partial
 from typing import Any, Callable, TypeVar
@@ -15,6 +16,34 @@ P = ParamSpec('P')
 R = TypeVar('R')
 
 
+class SubprocessException(Exception):
+    """A picklable exception to represent exceptions raised in subprocesses."""
+
+    def __init__(self, original_type, original_message, original_traceback) -> None:
+        self.original_type = original_type
+        self.original_message = original_message
+        self.original_traceback = original_traceback
+        super().__init__(f'{original_type}: {original_message}')
+
+    def __reduce__(self):
+        return (SubprocessException, (self.original_type, self.original_message, self.original_traceback))
+
+    def __str__(self):
+        return (f'Exception in subprocess:\n'
+                f'  Type: {self.original_type}\n'
+                f'  Message: {self.original_message}\n'
+                f'  {self.original_traceback}')
+
+
+def safe_callback(callback: Callable, *args, **kwargs) -> Any:
+    """Run a callback; catch and wrap any exceptions that might occur."""
+    try:
+        return callback(*args, **kwargs)
+    except Exception as e:
+        # NOTE: we do not want to pass the original exception because it might be unpicklable
+        raise SubprocessException(type(e).__name__, str(e), traceback.format_exc()) from None
+
+
 async def _run(executor: Any, callback: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
     if core.app.is_stopping:
         return  # type: ignore  # the assumption is that the user's code no longer cares about this value
@@ -37,7 +66,7 @@ async def cpu_bound(callback: Callable[P, R], *args: P.args, **kwargs: P.kwargs)
     It is encouraged to create static methods (or free functions) which get all the data as simple parameters (eg. no class/ui logic)
     and return the result (instead of writing it in class properties or global variables).
     """
-    return await _run(process_pool, callback, *args, **kwargs)
+    return await _run(process_pool, safe_callback, callback, *args, **kwargs)
 
 
 async def io_bound(callback: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:

tests/test_run.py

@@ -0,0 +1,72 @@
+import asyncio
+import time
+from typing import Awaitable, Generator
+
+import pytest
+
+from nicegui import app, run, ui
+from nicegui.testing import User
+
+
+@pytest.fixture(scope='module', autouse=True)
+def check_blocking_ui() -> Generator[None, None, None]:
+    """This fixture ensures that we see a warning if the UI is blocked for too long.
+
+    The warning would then automatically fail the test.
+    """
+    def configure() -> None:
+        loop = asyncio.get_running_loop()
+        loop.set_debug(True)
+        loop.slow_callback_duration = 0.02
+    app.on_startup(configure)
+    yield
+
+
+def delayed_hello() -> str:
+    """Test function that blocks for 1 second."""
+    time.sleep(1)
+    return 'hello'
+
+
+@pytest.mark.parametrize('func', [run.cpu_bound, run.io_bound])
+async def test_delayed_hello(user: User, func: Awaitable):
+    @ui.page('/')
+    async def index():
+        ui.label(await func(delayed_hello))
+
+    await user.open('/')
+    await user.should_see('hello')
+
+
+async def test_run_unpickable_exception_in_cpu_bound_callback(user: User):
+    class UnpicklableException(Exception):
+        def __reduce__(self):
+            raise NotImplementedError('This local object cannot be pickled')
+
+    def raise_unpicklable_exception():
+        raise UnpicklableException('test')
+
+    @ui.page('/')
+    async def index():
+        with pytest.raises(AttributeError, match="Can't pickle local object"):
+            ui.label(await run.cpu_bound(raise_unpicklable_exception))
+
+    await user.open('/')
+
+
+class ExceptionWithSuperParameter(Exception):
+    def __init__(self) -> None:
+        super().__init__('some parameter which does not appear in the custom exceptions init')
+
+
+def raise_exception_with_super_parameter():
+    raise ExceptionWithSuperParameter()
+
+
+async def test_run_cpu_bound_function_which_raises_problematic_exception(user: User):
+    @ui.page('/')
+    async def index():
+        with pytest.raises(run.SubprocessException, match='some parameter which does not appear in the custom exceptions init'):
+            ui.label(await run.cpu_bound(raise_exception_with_super_parameter))
+
+    await user.open('/')