From 6488344756c210e8eb20040bb5f71ff3033ad556 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Wed, 4 Oct 2023 15:00:50 -0700 Subject: [PATCH] improve coverage --- tests/tape_tests/test_ensemble.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/tests/tape_tests/test_ensemble.py b/tests/tape_tests/test_ensemble.py index ea6f0552..8f1801d3 100644 --- a/tests/tape_tests/test_ensemble.py +++ b/tests/tape_tests/test_ensemble.py @@ -499,10 +499,27 @@ def test_temporary_cols(parquet_ensemble): # try a sync ens._sync_tables() - # nobs_total should be removed + # nobs_total should be removed from object assert "nobs_total" not in ens._object_temp assert "nobs_total" not in ens._object.columns + # add a source column that we manually set as dirty, don't have a function + # that adds temporary source columns at the moment + ens.assign(f2=lambda x: x[ens._flux_col] ** 2, table="source") + ens._source_temp.append("f2") # manually append + + # prune object, object should be dirty + ens.prune(threshold=10) + + assert ens._object_dirty + + # try a sync + ens._sync_tables() + + # f2 should be removed from source + assert "f2" not in ens._source_temp + assert "f2" not in ens._source.columns + def test_dropna(parquet_ensemble): # Try passing in an unrecognized 'table' parameter and verify an exception is thrown @@ -597,10 +614,15 @@ def test_keep_zeros(parquet_ensemble): @pytest.mark.parametrize("by_band", [True, False]) -def test_calc_nobs(parquet_ensemble, by_band): +@pytest.mark.parametrize("know_divisions", [True, False]) +def test_calc_nobs(parquet_ensemble, by_band, know_divisions): ens = parquet_ensemble ens._object = ens._object.drop(["nobs_g", "nobs_r", "nobs_total"], axis=1) + if know_divisions: + ens._object = ens._object.reset_index().set_index(ens._id_col) + assert ens._object.known_divisions + ens.calc_nobs(by_band) lc = ens._object.loc[88472935274829959].compute()