diff --git a/plugins/ui/test/deephaven/ui/test_hooks.py b/plugins/ui/test/deephaven/ui/test_hooks.py index 40cdf1900..df9ed7e35 100644 --- a/plugins/ui/test/deephaven/ui/test_hooks.py +++ b/plugins/ui/test/deephaven/ui/test_hooks.py @@ -2,14 +2,16 @@ 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 @@ -17,10 +19,13 @@ def render_hook(fn: Callable): 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):