test: clean up connections

ibis/backends/athena/tests/conftest.py

@@ -118,4 +118,5 @@
 
 @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

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

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

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

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

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

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

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

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

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(

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

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

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:

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

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

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

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"<BackendTest {self.name()}>"
 

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