Skip to content

Commit

Permalink
test: clean up connections
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Feb 2, 2025
1 parent 622de09 commit 6b194a9
Show file tree
Hide file tree
Showing 18 changed files with 55 additions and 37 deletions.
3 changes: 2 additions & 1 deletion ibis/backends/athena/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 122 in ibis/backends/athena/tests/conftest.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/athena/tests/conftest.py#L121-L122

Added lines #L121 - L122 were not covered by tests
3 changes: 2 additions & 1 deletion ibis/backends/clickhouse/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
17 changes: 9 additions & 8 deletions ibis/backends/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand Down
5 changes: 3 additions & 2 deletions ibis/backends/datafusion/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 3 additions & 2 deletions ibis/backends/duckdb/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 2 additions & 3 deletions ibis/backends/flink/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 163 in ibis/backends/flink/tests/conftest.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/flink/tests/conftest.py#L162-L163

Added lines #L162 - L163 were not covered by tests


@pytest.fixture
Expand Down
3 changes: 2 additions & 1 deletion ibis/backends/impala/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 3 additions & 2 deletions ibis/backends/mssql/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
3 changes: 2 additions & 1 deletion ibis/backends/mysql/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 3 additions & 2 deletions ibis/backends/oracle/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 3 additions & 2 deletions ibis/backends/polars/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
3 changes: 2 additions & 1 deletion ibis/backends/postgres/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
8 changes: 3 additions & 5 deletions ibis/backends/pyspark/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,11 +398,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:
Expand Down
3 changes: 2 additions & 1 deletion ibis/backends/risingwave/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 3 additions & 2 deletions ibis/backends/snowflake/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 3 additions & 2 deletions ibis/backends/sqlite/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 6 additions & 0 deletions ibis/backends/tests/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()}>"

Expand Down
3 changes: 2 additions & 1 deletion ibis/backends/trino/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 6b194a9

Please sign in to comment.