Add tests for parametrized load_pmaps{,_as_df}

invisible_cities/io/pmaps_io_test.py

@@ -17,9 +17,13 @@
 from ..core.testing_utils import exactly
 from ..evm .pmaps         import S1
 from ..evm .pmaps         import S2
+from ..evm .pmaps         import PMap
 from ..evm .pmaps_test    import pmaps    as pmap_gen
 from .                    import pmaps_io as pmpio
 
+from typing import Generator
+from typing import Mapping
+
 
 pmaps_data = namedtuple("pmaps_data", """evt_numbers      peak_numbers
                                          evt_numbers_pmt  peak_numbers_pmt
@@ -201,6 +205,18 @@ def test_load_pmaps_as_df_lazy(KrMC_pmaps_filename, KrMC_pmaps_dfs):
         assert_dataframes_equal(df_lazy, df_eager)
 
 
+def test_load_pmaps_as_df_eager_lazy(KrMC_pmaps_filename):
+    eager = pmpio.load_pmaps_as_df(KrMC_pmaps_filename, lazy=False)
+    lazy  = pmpio.load_pmaps_as_df(KrMC_pmaps_filename, lazy=True )
+    assert len(eager) == 5
+    assert all(isinstance(item, pd.DataFrame) for item in eager)
+
+    assert isinstance(lazy, Generator)
+    element = next(lazy)
+    assert len(element   ) == 5
+    assert all(isinstance(item, pd.DataFrame) for item in element)
+
+
 def test_load_pmaps_as_df_eager_without_ipmt(KrMC_pmaps_without_ipmt_filename, KrMC_pmaps_without_ipmt_dfs):
     true_dfs = KrMC_pmaps_without_ipmt_dfs
     read_dfs = pmpio.load_pmaps_as_df_eager(KrMC_pmaps_without_ipmt_filename)
@@ -249,6 +265,20 @@ def test_load_pmaps_lazy(KrMC_pmaps_filename):
         assert_PMap_equality(pmap_lazy, pmaps_eager[evt])
 
 
+def test_load_pmaps_eager_lazy(KrMC_pmaps_filename):
+    eager = pmpio.load_pmaps(KrMC_pmaps_filename, lazy=False)
+    lazy  = pmpio.load_pmaps(KrMC_pmaps_filename, lazy=True )
+
+    assert isinstance(eager, Mapping)
+    element = next(iter(eager.values()))
+    assert isinstance(element, PMap)
+
+    assert isinstance(lazy, Generator)
+    element = next(lazy)
+    assert len(element) == 2 # event, pmap
+    assert isinstance(element[1], PMap)
+
+
 @mark.parametrize("signal_type", (S1, S2))
 def test_build_pmt_responses(KrMC_pmaps_dfs, signal_type):
     if   signal_type is S1:   df,  _, _, pmt_df,      _ = KrMC_pmaps_dfs