fix failing test

plugins/ui/test/deephaven/ui/test_hooks.py

@@ -2,25 +2,30 @@
 import unittest
 from operator import itemgetter
 from queue import Queue
+from time import sleep
 from typing import Callable
 from unittest.mock import Mock
 from .BaseTest import BaseTestCase
 
 LISTENER_TIMEOUT = 2.0
+QUEUE_TIMEOUT = 1.0
 
 
-def render_hook(fn: Callable):
+def render_hook(fn: Callable, queue=None):
     """
     Render a hook function and return the context, result, and a rerender function for updating it
 
     Args:
       fn: Callable:
         The function to render. Pass in a function with a hook call within it.
         Re-render will call the same function but with the new args passed in.
+      queue: Queue:
+        The queue to put items on. If not provided, a new queue will be created.
     """
     from deephaven.ui._internal.RenderContext import RenderContext
 
-    queue = Queue()
+    if queue is None:
+        queue = Queue()
 
     context = RenderContext(lambda x: queue.put(x), lambda x: queue.put(x))
 
@@ -42,6 +47,47 @@ def _rerender(*args, **kwargs):
     return return_dict
 
 
+class NotifyQueue(Queue):
+    """
+    A queue that notifies a function when an item is put on it
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._notify_fn = None
+
+    def put(self, item: object, block: bool = True, timeout: float = None) -> None:
+        """
+        Put an item on the queue and notify the function
+
+        Args:
+            item: The item to put on the queue
+            block: True if the call should block until the item is put on the queue
+            timeout: The time to wait for the item to be put on the queue
+
+        Returns:
+            None
+        """
+        super().put(item)
+        if self._notify_fn:
+            self._notify_fn(self)
+
+    def call_after_put(self, fn: Callable[["NotifyQueue"], None]) -> None:
+        """
+        Register a function to be called after an item is put on the queue
+
+        Args:
+            fn: The function to call after an item is put on the queue
+        """
+        self._notify_fn = fn
+
+    def unregister_notify(self) -> None:
+        """
+        Unregister the notify function
+        """
+        self._notify_fn = None
+
+
 class HooksTest(BaseTestCase):
     def test_state(self):
         from deephaven.ui.hooks import use_state
@@ -274,6 +320,36 @@ def _test_table_data(t=table):
 
         self.assertEqual(result, expected)
 
+    def verify_queue_has_size(self, queue: NotifyQueue, size: int):
+        """
+        Verify that the queue has the expected size in a multi-threaded context
+
+        Args:
+            queue: The queue to check
+            size: The expected size of the quexue
+
+        Returns:
+            None
+        """
+        event = threading.Event()
+
+        def check_size(q):
+            if q.qsize() == size:
+                event.set()
+
+        # call after each put in case the queue is not at the correct size yet
+        queue.call_after_put(check_size)
+
+        # call now in case the queue is (or was) already at the correct size
+        if check_size(queue):
+            queue.unregister_notify()
+            return
+
+        if not event.wait(timeout=QUEUE_TIMEOUT):
+            assert False, f"queue did not reach size {size}"
+
+        queue.unregister_notify()
+
     def test_swapping_table_data(self):
         from deephaven.ui.hooks import use_table_data
         from deephaven import new_table
@@ -292,7 +368,9 @@ def _test_table_data(t=table):
             result = use_table_data(t, sentinel="sentinel")
             return result
 
-        render_result = render_hook(_test_table_data)
+        queue = NotifyQueue()
+
+        render_result = render_hook(_test_table_data, queue)
 
         result, rerender = itemgetter("result", "rerender")(render_result)
 
@@ -313,10 +391,13 @@ def _test_table_data(t=table):
         self.verify_table_updated(table_writer, dynamic_table, (1, "Testing"))
 
         rerender(dynamic_table)
+        # the queue should have two items eventually, one for each set_state in _set_new_data in use_table_data
+        # this check is needed because the set_state calls come from the listener, which is called in a separate thread
+        # so the queue might not have the correct size immediately
+        self.verify_queue_has_size(queue, 2)
         result = rerender(dynamic_table)
 
         expected = {"Numbers": [1], "Words": ["Testing"]}
-
         self.assertEqual(result, expected)
 
     def test_column_data(self):