From 5b4407a31c0b745ad58d36e9c1afec812273cefb Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Tue, 13 Aug 2024 16:04:23 -0700
Subject: [PATCH] Fix race condition where the error box didn't always display
 (#788)

---
 mesop/examples/__init__.py                  |  1 +
 mesop/examples/event_handler_error.py       | 41 +++++++++++++++++++++
 mesop/tests/e2e/event_handler_error_test.ts | 37 +++++++++++++++++++
 mesop/web/src/shell/shell.ts                |  6 ++-
 4 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 mesop/examples/event_handler_error.py
 create mode 100644 mesop/tests/e2e/event_handler_error_test.ts

diff --git a/mesop/examples/__init__.py b/mesop/examples/__init__.py
index d7872aa4e..cfb2f12d2 100644
--- a/mesop/examples/__init__.py
+++ b/mesop/examples/__init__.py
@@ -16,6 +16,7 @@
 from mesop.examples import (
   error_no_stateclass_decorator as error_no_stateclass_decorator,
 )
+from mesop.examples import event_handler_error as event_handler_error
 from mesop.examples import focus_component as focus_component
 from mesop.examples import generator as generator
 from mesop.examples import grid as grid
diff --git a/mesop/examples/event_handler_error.py b/mesop/examples/event_handler_error.py
new file mode 100644
index 000000000..fe1acd653
--- /dev/null
+++ b/mesop/examples/event_handler_error.py
@@ -0,0 +1,41 @@
+import asyncio
+import time
+from typing import Generator
+
+import mesop as me
+
+
+@me.page(path="/event_handler_error")
+def page():
+  me.text("Hello, world!")
+  me.button(label="Regular event handler", on_click=on_click)
+  me.button(label="Generator event handler", on_click=on_click_generator)
+  me.button(label="Yield from event handler", on_click=on_click_yield_from)
+  me.button(
+    label="Async generator event handler", on_click=on_click_async_generator
+  )
+
+
+def on_click(e: me.ClickEvent):
+  raise Exception("Error in on_click.")
+
+
+def on_click_generator(e: me.ClickEvent):
+  yield
+  raise Exception("Error in on_click_generator.")
+
+
+def on_click_yield_from(e: me.ClickEvent) -> Generator[None, None, None]:
+  yield from a_generator()
+
+
+def a_generator():
+  yield
+  time.sleep(0.2)
+  raise Exception("Error in a_generator.")
+
+
+async def on_click_async_generator(e: me.ClickEvent):
+  await asyncio.sleep(0.2)
+  yield
+  raise Exception("Error in on_click_async_generator.")
diff --git a/mesop/tests/e2e/event_handler_error_test.ts b/mesop/tests/e2e/event_handler_error_test.ts
new file mode 100644
index 000000000..cce16ab88
--- /dev/null
+++ b/mesop/tests/e2e/event_handler_error_test.ts
@@ -0,0 +1,37 @@
+import {test, expect} from '@playwright/test';
+
+test('event handler error message is shown', async ({page}) => {
+  await page.goto('/event_handler_error');
+
+  // Regular event handler
+  await page.getByRole('button', {name: 'Regular event handler'}).click();
+  await expect(
+    page.getByText('Error in on_click.', {exact: true}),
+  ).toBeVisible();
+  await page.locator('button').filter({hasText: 'close'}).click();
+
+  // Generator event handler
+  await page
+    .getByRole('button', {name: 'Generator event handler', exact: true})
+    .click();
+  await expect(
+    page.getByText('Error in on_click_generator.', {exact: true}),
+  ).toBeVisible();
+  await page.locator('button').filter({hasText: 'close'}).click();
+
+  // Yield from event handler
+  await page.getByRole('button', {name: 'Yield from event handler'}).click();
+  await expect(
+    page.getByText('Error in a_generator.', {exact: true}),
+  ).toBeVisible();
+  await page.locator('button').filter({hasText: 'close'}).click();
+
+  // Async generator event handler
+  await page
+    .getByRole('button', {name: 'Async generator event handler'})
+    .click();
+  await expect(
+    page.getByText('Error in on_click_async_generator.', {exact: true}),
+  ).toBeVisible();
+  await page.locator('button').filter({hasText: 'close'}).click();
+});
diff --git a/mesop/web/src/shell/shell.ts b/mesop/web/src/shell/shell.ts
index 56c39db57..3ed1a20f6 100644
--- a/mesop/web/src/shell/shell.ts
+++ b/mesop/web/src/shell/shell.ts
@@ -97,6 +97,11 @@ export class Shell {
       {
         zone: this.zone,
         onRender: async (rootComponent, componentConfigs, jsModules) => {
+          // Make sure we clear the error *before* the async work, otherwise
+          // we can hit a weird race condition where the error is cleared
+          // right away before the user sees the error box.
+          this.error = undefined;
+
           // Import all JS modules concurrently
           await Promise.all(
             jsModules.map((modulePath) =>
@@ -113,7 +118,6 @@ export class Shell {
           if (componentConfigs.length) {
             this.componentConfigs = componentConfigs;
           }
-          this.error = undefined;
         },
         onCommand: (command) => {
           if (command.hasNavigate()) {