Skip to content

Commit

Permalink
add no client case
Browse files Browse the repository at this point in the history
  • Loading branch information
dougbrn committed Jan 31, 2024
1 parent 23e692f commit 3f35e03
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions tests/tape_tests/test_ensemble.py
Original file line number Diff line number Diff line change
Expand Up @@ -832,12 +832,19 @@ def test_persist(dask_client):
assert new_graph_size < old_graph_size


def test_sample_objects(parquet_ensemble_with_divisions):
@pytest.mark.parametrize(
"data_fixture",
[
"parquet_ensemble_with_divisions",
"parquet_ensemble_without_client",
],
)
def test_sample_objects(data_fixture, request):
"""
Test Ensemble.sample_objects
"""

ens = parquet_ensemble_with_divisions
ens = request.getfixturevalue(data_fixture)
ens.source.repartition(npartitions=10).update_ensemble()
ens.object.repartition(npartitions=5).update_ensemble()

Expand All @@ -853,7 +860,8 @@ def test_sample_objects(parquet_ensemble_with_divisions):
assert len(ens.object) == prior_obj_len
assert len(ens.source) == prior_src_len

ens.client.close()
if data_fixture == "parquet_ensemble_with_divisions":
ens.client.close()


def test_update_column_map(dask_client):
Expand Down

0 comments on commit 3f35e03

Please sign in to comment.