Fix unit tests

E3SM-Project · Feb 14, 2024 · e0a3660 · e0a3660
1 parent d76e971
commit e0a3660
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/e3sm_diags/driver/utils/regrid.py b/e3sm_diags/driver/utils/regrid.py
@@ -531,7 +531,7 @@ def _hybrid_to_plevs(
 
     # Vertical regriding sets the units to "mb", but the original units
     # should be preserved.
-    result[var_key].attrs["units"] = ds[var_key].units
+    result[var_key].attrs["units"] = ds[var_key].attrs["units"]
 
     return result
 

diff --git a/tests/e3sm_diags/driver/utils/test_regrid.py b/tests/e3sm_diags/driver/utils/test_regrid.py
@@ -334,7 +334,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_existing_z_bounds(self):
         # updating the arrays and attributes of data variables and coordinates.
         expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"])
         expected["so"].data[:] = np.nan
-        expected["so"].attrs["units"] = "mb"
+        expected["so"].attrs["units"] = "ppt"
         expected["lev"].attrs = {
             "axis": "Z",
             "coordinate": "vertical",
@@ -364,7 +364,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_generated_z_bounds(self):
         # updating the arrays and attributes of data variables and coordinates.
         expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"])
         expected["so"].data[:] = np.nan
-        expected["so"].attrs["units"] = "mb"
+        expected["so"].attrs["units"] = "ppt"
         expected["lev"].attrs = {
             "axis": "Z",
             "coordinate": "vertical",
@@ -393,7 +393,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_Pa_units(self):
         # updating the arrays and attributes of data variables and coordinates.
         expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"])
         expected["so"].data[:] = np.nan
-        expected["so"].attrs["units"] = "mb"
+        expected["so"].attrs["units"] = "ppt"
         expected["lev"].attrs = {
             "axis": "Z",
             "coordinate": "vertical",

diff --git a/tests/e3sm_diags/fixtures.py b/tests/e3sm_diags/fixtures.py
@@ -63,6 +63,7 @@ def generate_lev_dataset(
                 name="so",
                 data=np.ones((5, 4, 4, 4)),
                 coords={"time": time_decoded, "lev": lev, "lat": lat, "lon": lon},
+                attrs={"units": "ppt"},
             ),
         },
         coords={