Refactor process() testing & cleaning cube masking/filtering logic

test/test_um2netcdf.py

@@ -5,7 +5,217 @@
 import umpost.um2netcdf as um2nc
 
 import pytest
+import numpy as np
+
 import mule
+import mule.ff
+
+
+@pytest.fixture
+def z_sea_rho_data():
+    # data ripped from aiihca.paa1jan.subset: ff.level_dependent_constants.zsea_at_rho
+    # TODO: dtype is object, should it be float?
+    data = np.array([9.9982061118072, 49.998881525751194, 130.00023235363918,
+                     249.99833311211358, 410.00103476788956, 610.000486354252,
+                     850.0006133545584, 1130.0014157688088, 1449.9989681136456,
+                     1810.0011213557837, 2210.0000245285087, 2649.9996031151773,
+                     3129.9998571157903, 3650.000786530347, 4209.99846587549,
+                     4810.000746117935, 5449.999776290966, 6129.999481877941,
+                     6849.999862878861, 7610.000919293723, 8409.998725639172,
+                     9250.001132881924, 10130.00029005526, 11050.000122642545,
+                     12010.000630643768, 13010.001814058938, 14050.40014670717,
+                     15137.719781928794, 16284.973697054254, 17506.96881530842,
+                     18820.820244130424, 20246.59897992768, 21808.13663216417,
+                     23542.18357603375, 25520.960854349545, 27901.358260464756,
+                     31063.888598164976, 36081.76331548462, -1073741824.0], dtype=object)
+    return data
+
+
+@pytest.fixture
+def z_sea_theta_data():
+    # data ripped from aiihca.paa1jan.subset: ff.level_dependent_constants.zsea_at_theta
+    # TODO: dtype is object, should it be float?
+    data = np.array([0.0, 20.000337706971997, 80.00135082788799, 179.9991138793904,
+                     320.00147782819437, 500.00059170758476, 720.0003810009191,
+                     980.0008457081975, 1279.9980603460624, 1619.9998758812287,
+                     1999.9984413469813, 2420.001607710036, 2880.0015240036764,
+                     3379.9981902279037, 3919.9994573494323, 4500.001399884905,
+                     5120.000092350965, 5779.999460230968, 6479.999503524915,
+                     7220.000222232806, 8000.001616354641, 8819.999760407061,
+                     9679.998579873429, 10579.998074753737, 11519.998245047991,
+                     12499.999090756188, 13520.000611878331, 14580.799681536007,
+                     15694.639882321579, 16875.311437270288, 18138.62619334655,
+                     19503.01036943094, 20990.18759042441, 22626.081748420565,
+                     24458.285403646936, 26583.640230535515, 29219.080215877355,
+                     32908.69305496925, 39254.833576], dtype=object)
+    return data
+
+
+@pytest.fixture
+def mule_vars(z_sea_rho_data, z_sea_theta_data):
+    """Simulate mule variables from aiihca.paa1jan.subset data."""
+    d_lat = 1.25  # spacing manually copied from aiihca.paa1jan.subset file
+    d_lon = 1.875
+    return um2nc.MuleVars(um2nc.GRID_NEW_DYNAMICS, d_lat, d_lon, z_sea_rho_data, z_sea_theta_data)
+
+
+@pytest.fixture
+def air_temp_cube():
+    # copied from aiihca.paa1jan.subset file
+    return DummyCube(30204, "air_temperature")
+
+
+@pytest.fixture
+def precipitation_flux_cube():
+    # copied from aiihca.paa1jan.subset file
+    return DummyCube(5216, "precipitation_flux")
+
+
+@pytest.fixture
+def std_args():
+    # TODO: make args namedtuple?
+    args = mock.Mock()
+    args.nomask = False
+    args.nohist = False
+    args.nckind = 3
+    args.include_list = None
+    args.exclude_list = None
+    args.simple = False
+    args.verbose = False
+    return args
+
+
+@pytest.fixture
+def fake_in_path():
+    # use junk paths to protect against accidentally touching filesystems
+    return "/tmp-does-not-exist/fake_input_fields_file"
+
+
+@pytest.fixture
+def fake_out_path():
+    # use junk paths to protect against accidentally touching filesystems
+    return "/tmp-does-not-exist/fake_input_fields_file.nc"
+
+
+def test_process_without_masking(air_temp_cube, precipitation_flux_cube, mule_vars,
+                                 std_args, fake_in_path, fake_out_path):
+    """Attempts end-to-end test of process(), ignoring cubes requiring masking."""
+
+    # FIXME: this convoluted setup is a big code stench
+    #        use these tests to gradually refactor process()
+    # TODO: move towards a design where input & output I/O is extracted from process()
+    #       process()'s core should operate with *data only* args
+    with mock.patch("mule.load_umfile"):  # ignore m_load_umfile as process_mule_vars is mocked
+        with mock.patch("umpost.um2netcdf.process_mule_vars") as m_mule_vars:
+            m_mule_vars.return_value = mule_vars
+
+            with mock.patch("iris.load") as m_iris_load:
+                cubes = [air_temp_cube, precipitation_flux_cube]
+
+                for c in cubes:
+                    c.attributes = {um2nc.STASH: DummyStash(*split_item_code(c.item_code))}
+                    c.cell_methods = []
+
+                    # TODO: replace with DummyCubeWithCoords ()?
+                    #       or add coord fixtures?
+                    c.coord["latitude"] = 0.0  # FIXME
+                    c.coord["longitude"] = 0.0  # FIXME
+
+                m_iris_load.return_value = cubes
+
+                with mock.patch("iris.fileformats.netcdf.Saver") as m_saver:  # prevent I/O
+                    # mock `sman` var to prevent I/O
+                    m_sman = mock.Mock()
+                    m_saver().__enter__.return_value = m_sman
+
+                    # TODO: fix lat/lon & levels requires c.coord attribute
+                    #       use fixtures to add attrs & remove the patches?
+                    with mock.patch("umpost.um2netcdf.fix_latlon_coord") as m_coord:
+                        with mock.patch("umpost.um2netcdf.fix_level_coord") as m_level:
+                            with mock.patch("umpost.um2netcdf.apply_mask") as m_apply_mask:
+                                with mock.patch("umpost.um2netcdf.cubewrite") as m_cubewrite:
+                                    std_args.verbose = True  # test some warning branches
+                                    um2nc.process(fake_in_path, fake_out_path, std_args)
+
+                                    assert m_sman.update_global_attributes.called
+                                    assert m_saver.write.called is False  # write I/O prevented
+                                    assert m_coord.called
+                                    assert m_level.called
+                                    assert m_apply_mask.called is False
+                                    assert m_cubewrite.called  # real cubewrite() prevented
+                                    assert m_cubewrite.call_count == 1
+                                    assert m_cubewrite.call_args_list[0].args[0] == precipitation_flux_cube
+
+
+def test_process_all_cubes_filtered(air_temp_cube, mule_vars, std_args,
+                                    fake_in_path, fake_out_path):
+    """Ensure process() exists early if all cubes are removed in filtering."""
+    with mock.patch("mule.load_umfile"):  # ignore m_load_umfile as process_mule_vars is mocked
+        with mock.patch("umpost.um2netcdf.process_mule_vars") as m_mule_vars:
+            m_mule_vars.return_value = mule_vars
+
+            with mock.patch("iris.load") as m_iris_load:
+                section, item = split_item_code(air_temp_cube.item_code)
+                air_temp_cube.attributes = {um2nc.STASH: DummyStash(section, item)}
+                m_iris_load.return_value = [air_temp_cube]
+
+                with mock.patch("iris.fileformats.netcdf.Saver") as m_saver:  # prevent I/O
+                    m_sman = mock.Mock()
+                    m_saver().__enter__.return_value = m_sman
+
+                    um2nc.process(fake_in_path, fake_out_path, std_args)
+
+                    assert m_sman.update_global_attributes.called is False
+                    assert m_saver.write.called is False  # write I/O prevented
+
+
+def test_process_masking(air_temp_cube, precipitation_flux_cube,
+                         heaviside_uv_cube, heaviside_t_cube,
+                         mule_vars, std_args, fake_in_path, fake_out_path):
+    """Run process() with masking cubes."""
+    with mock.patch("mule.load_umfile"):  # ignore m_load_umfile as process_mule_vars is mocked
+        with mock.patch("umpost.um2netcdf.process_mule_vars") as m_mule_vars:
+            m_mule_vars.return_value = mule_vars
+
+            with mock.patch("iris.load") as m_iris_load:
+                # add cube requiring heaviside_t masking to enable both uv & t code branches
+                geo_potential_cube = DummyCube(30297, "geopotential_height")
+
+                cubes = [air_temp_cube, precipitation_flux_cube, geo_potential_cube,
+                         heaviside_uv_cube, heaviside_t_cube]
+
+                for c in cubes:
+                    attrs = {um2nc.STASH: DummyStash(*split_item_code(c.item_code))}
+                    c.attributes = attrs
+                    c.cell_methods = []
+
+                m_iris_load.return_value = cubes
+
+                with mock.patch("iris.fileformats.netcdf.Saver") as m_saver:  # prevent I/O
+                    m_sman = mock.Mock()
+                    m_saver().__enter__.return_value = m_sman
+
+                    # TODO: fix lat/lon & levels requires c.coord attributes
+                    #       use fixtures to add attrs & remove the patches?
+                    with mock.patch("umpost.um2netcdf.fix_latlon_coord") as m_coord:
+                        with mock.patch("umpost.um2netcdf.fix_level_coord") as m_level:
+                            with mock.patch("umpost.um2netcdf.apply_mask") as m_apply_mask:
+                                with mock.patch("umpost.um2netcdf.cubewrite") as m_cubewrite:
+                                    um2nc.process(fake_in_path, fake_out_path, std_args)
+
+                                    assert m_sman.update_global_attributes.called
+                                    assert m_sman.update_global_attributes.call_count == 2
+                                    assert m_saver.write.called is False  # write I/O prevented
+                                    assert m_coord.called
+                                    assert m_level.called
+                                    assert m_apply_mask.called
+                                    assert m_cubewrite.called  # real cubewrite() should be prevented
+                                    assert m_cubewrite.call_count == len(cubes)
+
+
+def split_item_code(item_code: int):
+    """Helper func: convert item code back to older section & item components."""
+    return item_code // 1000, item_code % 1000
 
 
 def test_get_eg_grid_type():
@@ -79,7 +289,7 @@ def test_stash_code_to_item_code_conversion():
     assert result == 30255
 
 
-@dataclass
+@dataclass(frozen=True)
 class DummyStash:
     """
     Partial Stash representation for testing.
@@ -127,12 +337,23 @@ def __init__(self, item_code, var_name=None, attributes=None, units=None):
         self.var_name = var_name or "unknown_var"
         self.attributes = attributes
         self.units = None or units
+        self.standard_name = None
+        self.long_name = None
+        self.coord = {}
+
+    def name(self):
+        # mimic iris API
+        return self.var_name
 
 
-def test_set_item_codes_fail_on_overwrite():
-    cubes = [DummyCube(1007, "fake_var")]
-    with pytest.raises(NotImplementedError):
-        um2nc.set_item_codes(cubes)
+def test_set_item_codes_avoid_overwrite():
+    item_code = 1007
+    item_code2 = 51006
+
+    cubes = [DummyCube(item_code, "fake_var"), DummyCube(item_code2, "fake_var2")]
+    um2nc.set_item_codes(cubes)
+    assert cubes[0].item_code == item_code
+    assert cubes[1].item_code == item_code2
 
 
 @pytest.fixture
@@ -150,6 +371,11 @@ def ta_plev_cube():
     return DummyCube(30294, "ta_plev")
 
 
+@pytest.fixture
+def heaviside_t_cube():
+    return DummyCube(30304, "heaviside_t")
+
+
 def test_check_pressure_level_masking_need_heaviside_uv(ua_plev_cube,
                                                         heaviside_uv_cube):
     cubes = [ua_plev_cube, heaviside_uv_cube]
@@ -171,8 +397,7 @@ def test_check_pressure_level_masking_missing_heaviside_uv(ua_plev_cube):
     assert heaviside_uv is None
 
 
-def test_check_pressure_level_masking_need_heaviside_t(ta_plev_cube):
-    heaviside_t_cube = DummyCube(30304)
+def test_check_pressure_level_masking_need_heaviside_t(ta_plev_cube, heaviside_t_cube):
     cubes = (ta_plev_cube, heaviside_t_cube)
 
     (need_heaviside_uv, heaviside_uv,