From f1bc699f8a80bad86607901aa5954b8ebb8f969c Mon Sep 17 00:00:00 2001
From: Panu Lahtinen <panu.lahtinen@fmi.fi>
Date: Thu, 7 Mar 2024 16:44:34 +0200
Subject: [PATCH] Refactor and test getting netcdf variable to xr.DataArray

---
 satpy/readers/netcdf_utils.py                 | 29 ++++++----
 satpy/tests/reader_tests/test_netcdf_utils.py | 57 +++++++++++++++++++
 2 files changed, 75 insertions(+), 11 deletions(-)

diff --git a/satpy/readers/netcdf_utils.py b/satpy/readers/netcdf_utils.py
index 560039c406..457c2ecaea 100644
--- a/satpy/readers/netcdf_utils.py
+++ b/satpy/readers/netcdf_utils.py
@@ -258,17 +258,7 @@ def collect_cache_vars(self, cache_var_size):
         cache_vars = self._collect_cache_var_names(cache_var_size)
         for var_name in cache_vars:
             v = self.file_content[var_name]
-            try:
-                attrs = v.attrs
-            except AttributeError:
-                attrs = v.__dict__
-            try:
-                arr = xr.DataArray(
-                    v[:], dims=v.dimensions, attrs=attrs, name=v.name)
-            except (ValueError, IndexError):
-                # Handle scalars for h5netcdf backend
-                arr = xr.DataArray(
-                    v.__array__(), dims=v.dimensions, attrs=attrs, name=v.name)
+            arr = get_data_as_xarray(v)
             self.cached_file_content[var_name] = arr
 
     def _collect_cache_var_names(self, cache_var_size):
@@ -384,6 +374,23 @@ def _compose_replacement_names(variable_name_replacements, var, variable_names):
                 variable_names.append(var.format(**{key: val}))
 
 
+def get_data_as_xarray(variable):
+    """Get data in variable as xr.DataArray."""
+    try:
+        attrs = variable.attrs
+    except AttributeError:
+        attrs = variable.__dict__
+    try:
+        arr = xr.DataArray(
+            variable[:], dims=variable.dimensions, attrs=attrs, name=variable.name)
+    except (ValueError, IndexError):
+        # Handle scalars for h5netcdf backend
+        arr = xr.DataArray(
+            variable.__array__(), dims=variable.dimensions, attrs=attrs, name=variable.name)
+
+    return arr
+
+
 class NetCDF4FsspecFileHandler(NetCDF4FileHandler):
     """NetCDF4 file handler using fsspec to read files remotely."""
 
diff --git a/satpy/tests/reader_tests/test_netcdf_utils.py b/satpy/tests/reader_tests/test_netcdf_utils.py
index 2d29288784..ad7b05c217 100644
--- a/satpy/tests/reader_tests/test_netcdf_utils.py
+++ b/satpy/tests/reader_tests/test_netcdf_utils.py
@@ -293,3 +293,60 @@ def test_use_h5netcdf_for_file_not_accessible_locally(self):
                 fh = NetCDF4FsspecFileHandler(fname, {}, {})
                 h5_file.assert_called_once()
                 assert fh._use_h5netcdf
+
+
+NC_ATTRS = {
+    "standard_name": "test_data",
+    "scale_factor": 0.01,
+    "add_offset": 0}
+
+def test_get_data_as_xarray_netcdf4():
+    """Test getting xr.DataArray from netcdf4 variable."""
+    import tempfile
+
+    import netCDF4 as nc
+    import numpy as np
+
+    from satpy.readers.netcdf_utils import get_data_as_xarray
+
+    data = np.array([1, 2, 3])
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create an empty HDF5
+        fname = os.path.join(tmpdir, "test.nc")
+        dset = nc.Dataset(fname, "w")
+        dset.createDimension("y", None)
+        var = dset.createVariable("test_data", "uint8", ("y",))
+        var[:] = data
+        var.setncatts(NC_ATTRS)
+        # Turn off automatic scale factor and offset handling
+        dset.set_auto_maskandscale(False)
+        res = get_data_as_xarray(var)
+        np.testing.assert_equal(res.data, data)
+        assert res.attrs == NC_ATTRS
+        dset.close()
+
+
+def test_get_data_as_xarray_h5netcdf():
+    """Test getting xr.DataArray from h5netcdf variable."""
+    import tempfile
+
+    import h5netcdf
+    import numpy as np
+
+    from satpy.readers.netcdf_utils import get_data_as_xarray
+
+    data = np.array([1, 2, 3])
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create an empty HDF5
+        fname = os.path.join(tmpdir, "test.nc")
+        with h5netcdf.File(fname, "w") as fid:
+            fid.dimensions = {"y": data.size}
+            var = fid.create_variable("test_data", ("y",), "uint8")
+            var[:] = data
+            for key in NC_ATTRS:
+                var.attrs[key] = NC_ATTRS[key]
+            res = get_data_as_xarray(var)
+            np.testing.assert_equal(res.data, data)
+            assert res.attrs == NC_ATTRS