Update pytest to generate a client fixture

.github/docker-compose.yaml

@@ -8,10 +8,12 @@ services:
         ports:
             - "8786:8786"
         environment:
-            EXTRA_CONDA_PACKAGES: "pandas>=1.3 numpy=1.20.2 -c conda-forge"
+            EXTRA_CONDA_PACKAGES: "pandas>=1.0.0"
     dask-worker:
         container_name: dask-worker
         image: daskdev/dask:latest
         command: dask-worker dask-scheduler:8786
         environment:
-            EXTRA_CONDA_PACKAGES: "pandas>=1.3 numpy=1.20.2 -c conda-forge"
+            EXTRA_CONDA_PACKAGES: "pandas>=1.0.0"
+        volumes:
+            - /tmp:/tmp

conftest.py

@@ -1,6 +1,6 @@
 import pytest
 
-pytest_plugins = ["distributed.utils_test", "tests.integration.fixtures"]
+pytest_plugins = ["tests.integration.fixtures"]
 
 
 def pytest_addoption(parser):

tests/integration/fixtures.py

@@ -6,7 +6,8 @@
 import pandas as pd
 import pytest
 from dask.datasets import timeseries
-from dask.distributed import Client
+from dask.distributed import Client, LocalCluster
+from dask.distributed.utils_test import loop  # noqa: F401
 from pandas.testing import assert_frame_equal
 
 try:
@@ -296,3 +297,15 @@ def setup_dask_client():
     os.getenv("DASK_SQL_TEST_SCHEDULER", None) is not None,
     reason="Can not run with external cluster",
 )
+
+
+@pytest.fixture()
+def cluster(loop):  # noqa: F811
+    with LocalCluster(loop=loop) as cluster:
+        yield cluster
+
+
+@pytest.fixture()
+def client(cluster):
+    with Client(cluster) as client:
+        yield client

tests/integration/test_cmd.py

@@ -1,4 +1,5 @@
 import pytest
+from dask import config as dask_config
 from mock import MagicMock, patch
 from prompt_toolkit.application import create_app_session
 from prompt_toolkit.input import create_pipe_input
@@ -103,8 +104,9 @@ def test_meta_commands(c, client, capsys):
         match="Timed out during handshake while "
         "connecting to tcp://localhost:8787 after 5 s",
     ):
-        client = _meta_commands("\\dsc localhost:8787", context=c, client=client)
-        assert client.scheduler.__dict__["addr"] == "localhost:8787"
+        with dask_config.set({"distributed.comm.timeouts.connect": 5}):
+            client = _meta_commands("\\dsc localhost:8787", context=c, client=client)
+            assert client.scheduler.__dict__["addr"] == "localhost:8787"
 
 
 def test_connection_info(c, client, capsys):

tests/integration/test_create.py

@@ -4,10 +4,8 @@
 from pandas.testing import assert_frame_equal
 
 import dask_sql
-from tests.integration.fixtures import skip_if_external_scheduler
 
 
-@skip_if_external_scheduler
 @pytest.mark.parametrize("gpu", [False, pytest.param(True, marks=pytest.mark.gpu)])
 def test_create_from_csv(c, df, temporary_data_file, gpu):
     df.to_csv(temporary_data_file, index=False)
@@ -66,7 +64,6 @@ def test_cluster_memory(client, c, df, gpu):
     assert_frame_equal(df, return_df)
 
 
-@skip_if_external_scheduler
 @pytest.mark.parametrize("gpu", [False, pytest.param(True, marks=pytest.mark.gpu)])
 def test_create_from_csv_persist(c, df, temporary_data_file, gpu):
     df.to_csv(temporary_data_file, index=False)
@@ -159,7 +156,6 @@ def test_create_from_query(c, df):
     assert_frame_equal(df, return_df)
 
 
-@skip_if_external_scheduler
 @pytest.mark.parametrize(
     "gpu",
     [

tests/integration/test_fugue.py

@@ -4,6 +4,7 @@
 from pandas.testing import assert_frame_equal
 
 from dask_sql import Context
+from tests.integration.fixtures import skip_if_external_scheduler
 
 fugue_sql = pytest.importorskip("fugue_sql")
 
@@ -38,6 +39,9 @@ def test_simple_statement():
     assert_frame_equal(return_df, pd.DataFrame({"a": [1], "b": ["world"]}))
 
 
+# TODO: Revisit fixing this on an independant cluster (without dask-sql) based on the
+# discussion in https://github.com/dask-contrib/dask-sql/issues/407
+@skip_if_external_scheduler
 def test_fsql():
     def assert_eq(df: pd.DataFrame) -> None:
         assert_frame_equal(df, pd.DataFrame({"a": [1]}))