retry

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"))

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)