From 94983a4f90a8a55db54a17ced31b7398c590132e Mon Sep 17 00:00:00 2001
From: Jared Lewis <jared@jared.kiwi.nz>
Date: Mon, 15 Jul 2024 09:09:45 +1000
Subject: [PATCH] test: Add test

---
 src/openmethane_prior/outputs.py | 21 +++++++++++----------
 src/openmethane_prior/raster.py  |  3 ++-
 tests/test_om_prior.py           | 11 +++++++++++
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/openmethane_prior/outputs.py b/src/openmethane_prior/outputs.py
index e9a27c0..68aac7f 100644
--- a/src/openmethane_prior/outputs.py
+++ b/src/openmethane_prior/outputs.py
@@ -22,14 +22,15 @@
 import numpy as np
 import numpy.typing as npt
 import xarray as xr
+
 from openmethane_prior.config import PriorConfig
 from openmethane_prior.layers import layer_names
 from openmethane_prior.utils import SECS_PER_YEAR
 
-coordNames = ["TSTEP", "LAY", "ROW", "COL"]
-requiredAttributes = {"units": "kg/m^2/s"}
-totalLayerName = "OCH4_TOTAL"
-totalLayerLongName = "total methane flux"
+COORD_NAMES = ["TSTEP", "LAY", "ROW", "COL"]
+REQUIRED_ATTRIBUTES = {"units": "kg/m^2/s"}
+TOTAL_LAYER_NAME = "OCH4_TOTAL"
+TOTAL_LAYER_LONG_NAME = "total methane flux"
 
 
 def convert_to_timescale(emission, cell_area):
@@ -77,9 +78,9 @@ def write_layer(
         # coerce to four dimensions if it's not
         for i in range(layer_data.ndim, 4):
             copy = np.expand_dims(copy, 0)  # should now have four dimensions
-        ds[layer_name] = (coordNames[:], copy)
+        ds[layer_name] = (COORD_NAMES[:], copy)
 
-    for k, v in requiredAttributes.items():
+    for k, v in REQUIRED_ATTRIBUTES.items():
         ds[layer_name].attrs[k] = v
     ds[layer_name].attrs["long_name"] = layer_name
 
@@ -118,8 +119,8 @@ def sum_layers(output_path: pathlib.Path):
             summed += ds[layerName].values  # it will broadcast time dimensions of 1 correctly
 
     if summed is not None:
-        ds[totalLayerName] = (["date", "LAY", *coordNames[-2:]], summed)
-        for k, v in requiredAttributes.items():
-            ds[totalLayerName].attrs[k] = v
-        ds[totalLayerName].attrs["long_name"] = totalLayerLongName
+        ds[TOTAL_LAYER_NAME] = (["date", "LAY", *COORD_NAMES[-2:]], summed)
+        for k, v in REQUIRED_ATTRIBUTES.items():
+            ds[TOTAL_LAYER_NAME].attrs[k] = v
+        ds[TOTAL_LAYER_NAME].attrs["long_name"] = TOTAL_LAYER_LONG_NAME
         ds.to_netcdf(output_path)
diff --git a/src/openmethane_prior/raster.py b/src/openmethane_prior/raster.py
index 8c13cda..ba1b98e 100644
--- a/src/openmethane_prior/raster.py
+++ b/src/openmethane_prior/raster.py
@@ -1,9 +1,10 @@
 import os
 
 import rasterio as rio
-from openmethane_prior.config import PriorConfig
 from rasterio.warp import Resampling, calculate_default_transform, reproject
 
+from openmethane_prior.config import PriorConfig
+
 
 def reproject_tiff(image, output, dst_crs="EPSG:4326", resampling="nearest", **kwargs):
     """Reprojects an image.
diff --git a/tests/test_om_prior.py b/tests/test_om_prior.py
index 7d324f5..a50d280 100644
--- a/tests/test_om_prior.py
+++ b/tests/test_om_prior.py
@@ -104,3 +104,14 @@ def test_compare_out_domain_with_cro_dot_files(output_domain, cro_xr, dot_xr):
 
     assert cro_xr.NCOLS == output_domain.COL.size
     assert cro_xr.NROWS == output_domain.ROW.size
+
+
+def test_required_attributes(output_domain):
+    assert output_domain.variables["OCH4_TOTAL"].attrs == {
+        "units": "kg/m^2/s",
+        "long_name": "OCH4_TOTAL",
+    }
+    assert output_domain.variables["OCH4_WETLANDS"].attrs == {
+        "units": "kg/m^2/s",
+        "long_name": "OCH4_WETLANDS",
+    }