diff --git a/tests/functional_tests/helpers.py b/tests/functional_tests/helpers.py index a95e84f1..a4068341 100644 --- a/tests/functional_tests/helpers.py +++ b/tests/functional_tests/helpers.py @@ -2,17 +2,18 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Sequence -from harlequin.app import Harlequin +from harlequin.app import Harlequin, ResultsFetched from harlequin.catalog import Catalog, CatalogItem from harlequin.components import ErrorModal from harlequin.components.code_editor import CodeEditor from harlequin.components.data_catalog.database_tree import DatabaseTree from harlequin.components.results_viewer import ResultsTable -from tests.waiting import wait_for, wait_for_value +from tests.waiting import wait_for, wait_for_messages, wait_for_value if TYPE_CHECKING: + from textual.message import Message from textual.pilot import Pilot from textual.widgets._tree import TreeNode @@ -42,6 +43,23 @@ async def wait_for_any_table(pilot: Pilot, app: Harlequin) -> ResultsTable: ) +async def wait_for_results( + pilot: Pilot, app: Harlequin, messages: Sequence[Message], *, count: int = 1 +) -> None: + """Wait until the app has finished handling the `count`th ResultsFetched. + + The message hook fires before the handler runs, and the handler awaits each + table's mount before it moves focus to the Results Viewer; `loading` clears + alongside that focus change. + """ + await wait_for_messages(pilot, messages, ResultsFetched, count=count) + await wait_for( + pilot, + lambda: not app.results_viewer.loading, + description="the Results Viewer to finish loading", + ) + + async def wait_for_error_modal(pilot: Pilot, app: Harlequin) -> ErrorModal: """The error modal on top of the stack, once the app has raised one.""" return await wait_for_value( diff --git a/tests/functional_tests/test_app.py b/tests/functional_tests/test_app.py index 1a1401ec..50a0a8c0 100644 --- a/tests/functional_tests/test_app.py +++ b/tests/functional_tests/test_app.py @@ -22,6 +22,7 @@ wait_for_any_table, wait_for_editor, wait_for_error_modal, + wait_for_results, ) from tests.waiting import wait_for, wait_for_messages, wait_for_value @@ -172,6 +173,7 @@ async def test_multiple_queries( pilot, messages, QuerySubmitted, exactly=True ) assert query_submitted_message.queries == ["select 1;"] + await wait_for_results(pilot, app, messages) table = await wait_for_any_table(pilot, app) assert table.row_count == table.source_row_count == 1 assert "hide-tabs" in app.results_viewer.classes @@ -192,7 +194,7 @@ async def test_multiple_queries( description="the Results Viewer to show a tab per query", ) assert "hide-tabs" not in app.results_viewer.classes - await wait_for_messages(pilot, messages, ResultsFetched, count=2) + await wait_for_results(pilot, app, messages, count=2) await pilot.wait_for_scheduled_animations() snap_results.append(await app_snapshot(app, "Both queries")) assert app.results_viewer.active == "result-1" @@ -232,7 +234,7 @@ async def test_single_query_terminated_with_semicolon( pilot, messages, QuerySubmitted, exactly=True ) assert query_submitted_message.queries == ["select 1;"] - await wait_for_any_table(pilot, app) + await wait_for_results(pilot, app, messages) assert app.results_viewer.tab_count == 1 editor.focus() @@ -245,7 +247,7 @@ async def test_single_query_terminated_with_semicolon( pilot, messages, QuerySubmitted, count=2, exactly=True ) assert query_submitted_message.queries == ["select 1;"] - await wait_for_workers(app) + await wait_for_results(pilot, app, messages, count=2) assert app.results_viewer.tab_count == 1 editor.focus() @@ -256,7 +258,7 @@ async def test_single_query_terminated_with_semicolon( pilot, messages, QuerySubmitted, count=3, exactly=True ) assert query_submitted_message.queries == ["select 1;"] - await wait_for_workers(app) + await wait_for_results(pilot, app, messages, count=3) assert app.results_viewer.tab_count == 1