update test_helpers.py

test/test_helpers/test_helpers.py

@@ -9,10 +9,13 @@ import os
 
 import mock
 import pytest
+import string
 
 from mlair.helpers import to_list, dict_to_xarray, float_round, remove_items, extract_value, select_from_dict
 from mlair.helpers import PyTestRegex
 from mlair.helpers import Logger, TimeTracking
+from mlair.helpers.helpers import is_xarray, convert2xrda, convert_1d_np2xr
+
 
 
 class TestToList:
@@ -312,3 +315,105 @@ class TestExtractValue:
             extract_value([1, 2, 3])
         assert "Trying to extract an encapsulated value from objects with more than a single entry is not supported " \
                "by this function." in e.value.args[0]
+
+
+class TestIsXarray:
+
+    @pytest.fixture
+    def custom_xr_data(self):
+        return xr.DataArray(np.array(range(5)))
+
+    def test_is_xarray_xr_input(self, custom_xr_data):
+        # data_array = xr.DataArray(np.array(range(5)))
+        assert is_xarray(custom_xr_data) is True
+        assert is_xarray(xr.Dataset({'test': custom_xr_data})) is True
+
+    def test_is_xarray_other_input(self, custom_xr_data):
+        assert is_xarray(1) is False
+        assert is_xarray(1.) is False
+        assert is_xarray([1, 2.]) is False
+        assert is_xarray([custom_xr_data]) is False
+
+
+class TestConvert2xrDa:
+
+    @pytest.fixture
+    def custom_1d_npdata(self):
+        return np.array(range(9))
+
+    @pytest.fixture()
+    def custom_2d_npdata(self, custom_1d_npdata):
+        return np.stack([custom_1d_npdata, 2 * custom_1d_npdata])
+
+    @pytest.fixture
+    def custom_xr_dataarray(self, custom_1d_npdata):
+        return xr.DataArray(custom_1d_npdata)
+
+    @pytest.fixture
+    def custom_xr_dataset(self, custom_xr_dataarray):
+        return xr.Dataset({'test_1': custom_xr_dataarray})
+
+    def test_convert2xrda_xrdata_in(self, custom_xr_dataarray, custom_xr_dataset):
+        assert (convert2xrda(custom_xr_dataarray) == custom_xr_dataarray).all()
+        assert (convert2xrda(custom_xr_dataset) == custom_xr_dataset).all()
+
+    def test_convert2xrda_npdata_in_nokwargs(self, custom_1d_npdata, custom_2d_npdata):
+        converted_data = convert2xrda(custom_1d_npdata)
+        assert isinstance(converted_data, xr.DataArray)
+        assert (converted_data.values == custom_1d_npdata).all()
+        assert converted_data.dims == ('dim_0',)
+        assert converted_data.dim_0.size == custom_1d_npdata.shape[0]
+
+        # Feed in a 2D-np.array without additional kwargs
+        converted_data = convert2xrda(custom_2d_npdata)
+        assert isinstance(converted_data, xr.DataArray)
+        assert (converted_data.values == custom_2d_npdata).all()
+        assert converted_data.dims == ('dim_0', 'dim_1')
+        assert converted_data.dim_0.size == custom_2d_npdata.shape[0]
+        assert converted_data.dim_1.size == custom_2d_npdata.shape[1]
+
+    def test_convert2xrda_npdata_in_nokwargs_default_true(self, custom_1d_npdata, custom_2d_npdata):
+        converted_data = convert2xrda(custom_1d_npdata, use_1d_default=True)
+        assert isinstance(converted_data, xr.DataArray)
+        assert (converted_data.values == custom_1d_npdata).all()
+        assert converted_data.dims == ('points',)
+        assert converted_data.points.size == custom_1d_npdata.shape[0]
+
+        # Feed in a 2D-np.array without additional kwargs
+        with pytest.raises(ValueError) as e:
+            converted_data = convert2xrda(custom_2d_npdata, use_1d_default=True)
+        assert "different number of dimensions on data and dims: 2 vs 1" in e.value.args[0]
+
+    @pytest.mark.parametrize("use_1d_default", (False, True))
+    def test_convert2xrda_npdata_in_kwargs(self, custom_1d_npdata, custom_2d_npdata, use_1d_default):
+        converted_data = convert2xrda(custom_1d_npdata, use_1d_default=use_1d_default, dims='other_points')
+        assert isinstance(converted_data, xr.DataArray)
+        assert (converted_data.values == custom_1d_npdata).all()
+        assert converted_data.dims == ('other_points',)
+        assert converted_data.other_points.size == custom_1d_npdata.shape[0]
+
+        # Feed in a 2D-np.array with correct additional kwargs
+        converted_data = convert2xrda(custom_2d_npdata, use_1d_default=use_1d_default,
+                                      dims=['test_dim_0', 'test_dim_1'],
+                                      coords={'test_dim_0': list(string.ascii_lowercase[:custom_2d_npdata.shape[0]]),
+                                              'test_dim_1': list(string.ascii_lowercase[:custom_2d_npdata.shape[1]]),
+                                              },
+                                      )
+        assert isinstance(converted_data, xr.DataArray)
+        assert (converted_data.values == custom_2d_npdata).all()
+        assert converted_data.dims == ('test_dim_0', 'test_dim_1')
+        assert (converted_data.coords['test_dim_0'].values == np.array(['a', 'b'])).all()
+
+    @pytest.mark.parametrize("scalar", (1, 2.))
+    def test_convert2xrda_int_float_in_nokwargs_default_true(self, scalar):
+        converted_data = convert2xrda(scalar, use_1d_default=True)
+        assert isinstance(converted_data, xr.DataArray)
+        assert converted_data.values == np.array([scalar])
+        assert converted_data.dims == ('points',)
+
+    @pytest.mark.parametrize("wrong_input", ({1: 'b'}, [1], 'abc'))
+    def test_convert2xrda_wrong_type_in_default_true_nokwargs(self, wrong_input):
+        with pytest.raises(TypeError) as e:
+            converted_data = convert2xrda(wrong_input, use_1d_default=True)
+        assert f"`arr' must be arry-like, int or float. But is of type {type(wrong_input)}" in e.value.args[0]
+