From 1e761a55facc8aff713a0fbc784618a89e5b7b38 Mon Sep 17 00:00:00 2001 From: Colin Ho Date: Fri, 23 Feb 2024 12:57:39 -0800 Subject: [PATCH] retry --- tests/integration/sql/conftest.py | 19 +++++++++++++++++++ tests/integration/sql/test_trino.py | 9 ++++----- 2 files changed, 23 insertions(+), 5 deletions(-) create mode 100644 tests/integration/sql/conftest.py diff --git a/tests/integration/sql/conftest.py b/tests/integration/sql/conftest.py new file mode 100644 index 0000000000..ba15c2e42e --- /dev/null +++ b/tests/integration/sql/conftest.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import pytest +import sqlalchemy +import tenacity + +TRINO_URL = "trino://user@localhost:8080/tpch" + + +@tenacity.retry( + stop=tenacity.stop_after_delay(60), + retry=tenacity.retry_if_exception_type(sqlalchemy.exc.DBAPIError), + wait=tenacity.wait_fixed(5), + reraise=True, +) +@pytest.fixture(scope="session") +def check_db_server_initialized() -> None: + with sqlalchemy.create_engine(TRINO_URL).connect() as conn: + conn.execute(sqlalchemy.text("SELECT 1")) diff --git a/tests/integration/sql/test_trino.py b/tests/integration/sql/test_trino.py index 53df3e564b..932d0ed8ff 100644 --- a/tests/integration/sql/test_trino.py +++ b/tests/integration/sql/test_trino.py @@ -4,12 +4,11 @@ import pytest import daft - -URL = "trino://user@localhost:8080/tpch" +from tests.integration.sql.conftest import TRINO_URL @pytest.mark.integration() -def test_trino_create_dataframe_ok() -> None: - df = daft.read_sql("SELECT * FROM tpch.sf1.nation", URL) - pd_df = pd.read_sql("SELECT * FROM tpch.sf1.nation", URL) +def test_trino_create_dataframe_ok(check_db_server_initialized) -> None: + df = daft.read_sql("SELECT * FROM tpch.sf1.nation", TRINO_URL) + pd_df = pd.read_sql("SELECT * FROM tpch.sf1.nation", TRINO_URL) assert df.to_pandas().equals(pd_df)