diff --git a/ibis/backends/athena/tests/conftest.py b/ibis/backends/athena/tests/conftest.py index bf9c0014ba43..a6c9475b9218 100644 --- a/ibis/backends/athena/tests/conftest.py +++ b/ibis/backends/athena/tests/conftest.py @@ -118,4 +118,5 @@ def awards_players(self): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection diff --git a/ibis/backends/clickhouse/tests/conftest.py b/ibis/backends/clickhouse/tests/conftest.py index a2880a7e52e2..30ed83107fa6 100644 --- a/ibis/backends/clickhouse/tests/conftest.py +++ b/ibis/backends/clickhouse/tests/conftest.py @@ -179,7 +179,8 @@ def add_catalog_and_schema(node): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="session") diff --git a/ibis/backends/conftest.py b/ibis/backends/conftest.py index c1ca256bc638..98145e68dae1 100644 --- a/ibis/backends/conftest.py +++ b/ibis/backends/conftest.py @@ -401,7 +401,8 @@ def backend(request, data_dir, tmp_path_factory, worker_id) -> BackendTest: """Return an instance of BackendTest, loaded with data.""" cls = _get_backend_conf(request.param) - return cls.load_data(data_dir, tmp_path_factory, worker_id) + with cls.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be @pytest.fixture(scope="session") @@ -414,7 +415,8 @@ def con(backend): def backend_no_data(request, data_dir, tmp_path_factory, worker_id): """Return an instance of BackendTest, with no data loaded.""" cls = _get_backend_conf(request.param) - return cls(data_dir=data_dir, tmpdir=tmp_path_factory, worker_id=worker_id) + with cls(data_dir=data_dir, tmpdir=tmp_path_factory, worker_id=worker_id) as be: + yield be @pytest.fixture(scope="session") @@ -464,7 +466,8 @@ def _setup_backend(request, data_dir, tmp_path_factory, worker_id): ) def ddl_backend(request, data_dir, tmp_path_factory, worker_id): """Set up the backends that are SQL-based.""" - return _setup_backend(request, data_dir, tmp_path_factory, worker_id) + with _setup_backend(request, data_dir, tmp_path_factory, worker_id) as be: + yield be @pytest.fixture(scope="session") @@ -473,14 +476,12 @@ def ddl_con(ddl_backend): return ddl_backend.connection -@pytest.fixture( - params=_get_backends_to_test(keep=("pyspark",)), - scope="session", -) +@pytest.fixture(params=_get_backends_to_test(keep=("pyspark",)), scope="session") def udf_backend(request, data_dir, tmp_path_factory, worker_id): """Runs the UDF-supporting backends.""" cls = _get_backend_conf(request.param) - return cls.load_data(data_dir, tmp_path_factory, worker_id) + with cls.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be @pytest.fixture(scope="session") diff --git a/ibis/backends/datafusion/tests/conftest.py b/ibis/backends/datafusion/tests/conftest.py index 5dbdadecb907..ee7f269a7efe 100644 --- a/ibis/backends/datafusion/tests/conftest.py +++ b/ibis/backends/datafusion/tests/conftest.py @@ -74,8 +74,9 @@ def add_catalog_and_schema(node): @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="session") diff --git a/ibis/backends/duckdb/tests/conftest.py b/ibis/backends/duckdb/tests/conftest.py index 629cfee2c9d1..1f517412595c 100644 --- a/ibis/backends/duckdb/tests/conftest.py +++ b/ibis/backends/duckdb/tests/conftest.py @@ -138,8 +138,9 @@ def add_catalog_and_schema(node): @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="session") diff --git a/ibis/backends/flink/tests/conftest.py b/ibis/backends/flink/tests/conftest.py index 9376c25dd70f..5fae4ed11e0a 100644 --- a/ibis/backends/flink/tests/conftest.py +++ b/ibis/backends/flink/tests/conftest.py @@ -159,9 +159,8 @@ def simple_table(simple_schema): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConfForStreaming.load_data( - data_dir, tmp_path_factory, worker_id - ).connection + with TestConfForStreaming.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture diff --git a/ibis/backends/impala/tests/conftest.py b/ibis/backends/impala/tests/conftest.py index 0611c313da65..922e7063898f 100644 --- a/ibis/backends/impala/tests/conftest.py +++ b/ibis/backends/impala/tests/conftest.py @@ -191,7 +191,8 @@ def test_data_dir(env): @pytest.fixture(scope="session") def backend(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id) + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be @pytest.fixture(scope="module") diff --git a/ibis/backends/mssql/tests/conftest.py b/ibis/backends/mssql/tests/conftest.py index 4dd86ed760da..90108eda3c0a 100644 --- a/ibis/backends/mssql/tests/conftest.py +++ b/ibis/backends/mssql/tests/conftest.py @@ -60,5 +60,6 @@ def connect(*, tmpdir, worker_id, **kw): @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection diff --git a/ibis/backends/mysql/tests/conftest.py b/ibis/backends/mysql/tests/conftest.py index b8e4f5734d34..9a0db913e8a2 100644 --- a/ibis/backends/mysql/tests/conftest.py +++ b/ibis/backends/mysql/tests/conftest.py @@ -76,4 +76,5 @@ def connect(*, tmpdir, worker_id, **kw): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection diff --git a/ibis/backends/oracle/tests/conftest.py b/ibis/backends/oracle/tests/conftest.py index adaa86c780f4..010e80ec492f 100644 --- a/ibis/backends/oracle/tests/conftest.py +++ b/ibis/backends/oracle/tests/conftest.py @@ -133,8 +133,9 @@ def format_table(name: str) -> str: @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection def init_oracle_database( diff --git a/ibis/backends/polars/tests/conftest.py b/ibis/backends/polars/tests/conftest.py index 8aec5127e305..5033fc25d813 100644 --- a/ibis/backends/polars/tests/conftest.py +++ b/ibis/backends/polars/tests/conftest.py @@ -44,8 +44,9 @@ def assert_series_equal(cls, left, right, *args, **kwargs) -> None: @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="session") diff --git a/ibis/backends/postgres/tests/conftest.py b/ibis/backends/postgres/tests/conftest.py index a0ec63dfae20..e465c071a087 100644 --- a/ibis/backends/postgres/tests/conftest.py +++ b/ibis/backends/postgres/tests/conftest.py @@ -70,7 +70,8 @@ def connect(*, tmpdir, worker_id, **kw): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="module") diff --git a/ibis/backends/pyspark/tests/conftest.py b/ibis/backends/pyspark/tests/conftest.py index 65df70f78142..fbf8ae7dff2d 100644 --- a/ibis/backends/pyspark/tests/conftest.py +++ b/ibis/backends/pyspark/tests/conftest.py @@ -234,6 +234,9 @@ class TestConf(BaseSparkTestConf, ServiceBackendTest): data_volume = "/data" service_name = "spark-connect" + def __exit__(self, *args, **kwargs): + pass + @property def parquet_dir(self) -> str: return self.data_volume @@ -398,11 +401,9 @@ def write_to_memory(self, expr, table_name): @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - backend_test = TestConf.load_data(data_dir, tmp_path_factory, worker_id) - con = backend_test.connection - - return con +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection class IbisWindow: diff --git a/ibis/backends/risingwave/tests/conftest.py b/ibis/backends/risingwave/tests/conftest.py index 06940a6573b6..06917b45d6dd 100644 --- a/ibis/backends/risingwave/tests/conftest.py +++ b/ibis/backends/risingwave/tests/conftest.py @@ -69,7 +69,8 @@ def connect(*, tmpdir, worker_id, port: int | None = None, **kw): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection @pytest.fixture(scope="module") diff --git a/ibis/backends/snowflake/tests/conftest.py b/ibis/backends/snowflake/tests/conftest.py index 8b309d71d7db..17e72007573c 100644 --- a/ibis/backends/snowflake/tests/conftest.py +++ b/ibis/backends/snowflake/tests/conftest.py @@ -215,5 +215,6 @@ def connect(*, tmpdir, worker_id, **kw) -> BaseBackend: @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection diff --git a/ibis/backends/sqlite/tests/conftest.py b/ibis/backends/sqlite/tests/conftest.py index 7268fb802042..32e1633d006b 100644 --- a/ibis/backends/sqlite/tests/conftest.py +++ b/ibis/backends/sqlite/tests/conftest.py @@ -55,5 +55,6 @@ def functional_alltypes(self) -> ir.Table: @pytest.fixture(scope="session") -def con(data_dir, tmp_path_factory, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +def con(tmp_path_factory, data_dir, worker_id): + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection diff --git a/ibis/backends/tests/base.py b/ibis/backends/tests/base.py index eed34fe4a473..a254b4c0d21f 100644 --- a/ibis/backends/tests/base.py +++ b/ibis/backends/tests/base.py @@ -101,6 +101,12 @@ def __init__(self, *, data_dir: Path, tmpdir, worker_id, **kw) -> None: self.data_dir = data_dir self.script_dir = data_dir.parent / "schema" + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.connection.disconnect() + def __str__(self): return f"" diff --git a/ibis/backends/trino/tests/conftest.py b/ibis/backends/trino/tests/conftest.py index 3d9610d5f2d7..633465f8b952 100644 --- a/ibis/backends/trino/tests/conftest.py +++ b/ibis/backends/trino/tests/conftest.py @@ -163,7 +163,8 @@ def awards_players(self): @pytest.fixture(scope="session") def con(tmp_path_factory, data_dir, worker_id): - return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection + with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be: + yield be.connection def generate_tpc_tables(suite_name, *, data_dir):