full frame testing

tests/conftest.py

@@ -14,6 +14,10 @@
 from typing import Any, Callable
 import git
 
+
+from pandas._testing import assert_dict_equal
+from pandas.testing import assert_frame_equal
+
 import pytest
 
 from generic_parser import DotDict
@@ -87,6 +91,17 @@ def to_string(val: Any):
     return to_string
 
 
+def assert_tfsdataframe_equal(df1, df2, **kwargs):
+    """ Wrapper to compare two TfsDataFrames with 
+    `assert_frame_equal` for the data and `assert_dict_equal` for the headers. 
+    
+    The `kwargs` are passed to `assert_frame_equal`, with the exception of `compare_keys`.
+    """
+    compare_keys = kwargs.get('compare_keys', True)
+    assert_frame_equal(df1, df2, **kwargs)
+    assert_dict_equal(df1.headers, df2.headers, compare_keys=compare_keys)
+
+
 # Model fixtures from /inputs/models -------------------------------------------
 # Hint: Before adding 25cm models, update files (see inj model folders, jdilly 2021)
 

tests/unit/test_kmod_averaging.py

@@ -1,21 +1,20 @@
-from collections.abc import Sequence
 import logging
-from pathlib import Path
 import shutil
+from collections.abc import Sequence
+from pathlib import Path
 
-import pandas.testing as pdt
 import pytest
 import tfs
 
 from omc3.optics_measurements.constants import BEAM, BEAM_DIR, BETA, NAME
+from omc3.plotting.plot_kmod_results import PARAM_BETA, PARAM_BETABEAT, PARAM_WAIST
 from omc3.scripts.kmod_average import (
     AVERAGED_BETASTAR_FILENAME,
     AVERAGED_BPM_FILENAME,
     EXT,
     average_kmod_results,
 )
-from omc3.plotting.plot_kmod_results import PARAM_BETA, PARAM_BETABEAT, PARAM_WAIST
-from tests.conftest import INPUTS, ids_str
+from tests.conftest import INPUTS, assert_tfsdataframe_equal, ids_str
 
 KMOD_INPUT_DIR = INPUTS / "kmod"
 REFERENCE_DIR = KMOD_INPUT_DIR / "references"
@@ -43,7 +42,7 @@ def test_kmod_averaging(tmp_path, ip, n_files):
     for out_name in get_all_tfs_filenames(ip, beta[0]):
         out_file = tfs.read(tmp_path / out_name)
         ref_file = tfs.read(ref_output_dir / out_name)
-        pdt.assert_frame_equal(out_file, ref_file, check_like=True)
+        assert_tfsdataframe_equal(out_file, ref_file, check_like=True)
 
 
 @pytest.mark.extended
@@ -80,7 +79,7 @@ def test_kmod_averaging_single_beam(tmp_path, beam, caplog):
         ref_file = tfs.read(ref_output_dir / out_name)
         if BEAM in ref_file.columns:
             ref_file = ref_file.loc[ref_file[BEAM] == beam, :].reset_index(drop=True)
-        pdt.assert_frame_equal(out_file, ref_file, check_like=True)
+        assert_tfsdataframe_equal(out_file, ref_file, check_like=True)
 
 
 # Helper ---

tests/unit/test_kmod_import.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 
-import pandas.testing as pdt
 import pytest
 import tfs
 
@@ -12,6 +11,7 @@
     EXT,
 )
 from omc3.scripts.kmod_import import import_kmod_data
+from tests.conftest import assert_tfsdataframe_equal, ids_str
 from tests.unit.test_kmod_averaging import (
     KMOD_INPUT_DIR,
     get_betastar_model,
@@ -26,6 +26,9 @@
 B1_RESULTS_OUTPUTS = KMOD_INPUT_DIR / "b1_imported"
 B2_RESULTS_OUTPUTS = KMOD_INPUT_DIR / "b2_imported"
 
+
+# Tests ---
+
 @pytest.mark.basic
 @pytest.mark.parametrize('beam', [1, 2])
 def test_kmod_import_averaged_folder_beam(tmp_path, beam):
@@ -48,13 +51,13 @@ def test_kmod_import_averaged_folder_beam(tmp_path, beam):
             beta_out = tfs.read(tmp_path / ref_path.name)
 
             # column order might have changed, but that's okay -> check_like=True
-            pdt.assert_frame_equal(beta_ref, beta_out, check_like=True)
+            assert_tfsdataframe_equal(beta_ref, beta_out, check_like=True)
 
 
 @pytest.mark.extended
 @pytest.mark.parametrize('beam', [1, 2])
 @pytest.mark.parametrize('files', ["bpm", "betastar", "bpm-betastar"])
-@pytest.mark.parametrize('read', [True, False])
+@pytest.mark.parametrize('read', [True, False], ids=ids_str("read{}"))
 def test_kmod_import_files_beam(tmp_path, beam, files, read):
     model = get_model_path(beam)
     beta = get_betastar_model(beam, ip=1)[0]
@@ -90,8 +93,10 @@ def test_kmod_import_files_beam(tmp_path, beam, files, read):
             beta_out = tfs.read(tmp_path / ref_path.name)
 
             # column order might have changed, but that's okay -> check_like=True
-            pdt.assert_frame_equal(beta_ref, beta_out, check_like=True)
+            assert_tfsdataframe_equal(beta_ref, beta_out, check_like=True)
+
 
+# Helper ---
 
 def _assert_correct_files_are_present(outputdir: Path, which: str = "bpm-betastar") -> None:
     """Simply checks the expected converted files are present in the outputdir"""

tests/unit/test_kmod_importer.py

@@ -1,18 +1,29 @@
-import pandas.testing as pdt
 import pytest
 import tfs
 
 from omc3.kmod_importer import AVERAGE_DIR, import_kmod_results
 from omc3.optics_measurements.constants import EXT
+from tests.conftest import assert_tfsdataframe_equal
 from tests.unit.test_kmod_averaging import (
     get_all_tfs_filenames as _get_averaged_filenames,
+)
+from tests.unit.test_kmod_averaging import (
     get_betastar_model,
     get_measurement_dir,
     get_reference_dir,
 )
-from tests.unit.test_kmod_import import get_model_path, _get_bpm_reference_path, _get_betastar_reference_path 
-from tests.unit.test_kmod_lumi_imbalance import REFERENCE_DIR, _get_effbetas_filename as _get_lumi_filename
+from tests.unit.test_kmod_import import (
+    _get_betastar_reference_path,
+    _get_bpm_reference_path,
+    get_model_path,
+)
+from tests.unit.test_kmod_lumi_imbalance import REFERENCE_DIR
+from tests.unit.test_kmod_lumi_imbalance import (
+    _get_effbetas_filename as _get_lumi_filename,
+)
+
 
+# Tests ---
 
 @pytest.mark.basic
 @pytest.mark.parametrize('beam', [1, 2])
@@ -42,17 +53,17 @@ def test_full_kmod_import_beam(tmp_path, beam):
         for out_name in _get_averaged_filenames(ip, beta=beta):
             out_file = tfs.read(average_dir / out_name)
             ref_file = tfs.read(get_reference_dir(ip, n_files=2) / out_name)
-            pdt.assert_frame_equal(out_file, ref_file, check_like=True)
+            assert_tfsdataframe_equal(out_file, ref_file, check_like=True)
 
     # lumi --
     eff_betas = tfs.read(average_dir / _get_lumi_filename(beta))
     eff_betas_ref = tfs.read(REFERENCE_DIR / _get_lumi_filename(beta))
-    pdt.assert_frame_equal(eff_betas_ref, eff_betas, check_like=True)
+    assert_tfsdataframe_equal(eff_betas_ref, eff_betas, check_like=True)
 
     # import --
     for plane in "xy":
         for ref_path in (_get_bpm_reference_path(beam, plane), _get_betastar_reference_path(beam, plane)):
             ref_file = tfs.read(ref_path)
             out_file = tfs.read(tmp_path / ref_path.name)
-            pdt.assert_frame_equal(ref_file, out_file, check_like=True)
+            assert_tfsdataframe_equal(ref_file, out_file, check_like=True)
 

tests/unit/test_kmod_lumi_imbalance.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-import pandas.testing as pdt
+import pytest
 import tfs
 
 from omc3.optics_measurements.constants import (
@@ -9,10 +9,12 @@
     EXT,
 )
 from omc3.scripts.kmod_lumi_imbalance import calculate_lumi_imbalance
+from tests.conftest import assert_tfsdataframe_equal
 from tests.unit.test_kmod_averaging import REFERENCE_DIR, get_reference_dir
 
 # Tests ---
 
+@pytest.mark.basic
 def test_kmod_lumi_imbalance(tmp_path):
     beta = 0.22
     path_beta_ip1 = _get_input_path(1, beta)
@@ -22,7 +24,7 @@ def test_kmod_lumi_imbalance(tmp_path):
 
     eff_betas = tfs.read(tmp_path / _get_effbetas_filename(beta))
     eff_betas_ref = tfs.read(REFERENCE_DIR / _get_effbetas_filename(beta))
-    pdt.assert_frame_equal(eff_betas_ref, eff_betas, check_like=True)
+    assert_tfsdataframe_equal(eff_betas_ref, eff_betas, check_like=True)
 
 
 # Helper ---