Updating unit tests

ajsockol · Jul 17, 2024 · add4b4d · add4b4d
1 parent 9dc4b44
commit add4b4d
Showing 1 changed file with 14 additions and 26 deletions.
diff --git a/tests/utils/test_data_utils.py b/tests/utils/test_data_utils.py
@@ -566,6 +566,7 @@ def test_convert_2d_to_1d():
     ds = xr.Dataset({
         'var': (('time', 'level'), data)
     }, coords={'time': [0, 1, 2], 'level': [10, 20]})
+    ds['level'].attrs['units'] = 'm'
 
     # Run the function
     result = convert_2d_to_1d(ds, parse='level')
@@ -576,26 +577,6 @@ def test_convert_2d_to_1d():
     np.testing.assert_array_equal(result['var_level_0'].values, [1, 3, 5])
     np.testing.assert_array_equal(result['var_level_1'].values, [2, 4, 6])
 
-    # Create a sample dataset
-    data = np.array([[1], [3], [5]])
-    ds = xr.Dataset({
-        'var': (('time', 'level'), data)
-    }, coords={'time': [0, 1, 2], 'level': [10]})
-
-    # Run the function with keep_name_if_one=True
-    result = convert_2d_to_1d(ds, parse='level', keep_name_if_one=True)
-
-    # Check the results
-    assert 'var' in result
-    np.testing.assert_array_equal(result['var'].values, [1, 3, 5])
-
-    # Create a sample dataset
-    data = np.array([[1, 2], [3, 4], [5, 6]])
-    ds = xr.Dataset({
-        'var': (('time', 'level'), data)
-    }, coords={'time': [0, 1, 2], 'level': [10, 20]})
-    ds['level'].attrs['units'] = 'm'
-
     # Run the function with use_dim_value_in_name=True
     result = convert_2d_to_1d(ds, parse='level', use_dim_value_in_name=True)
 
@@ -605,12 +586,6 @@ def test_convert_2d_to_1d():
     np.testing.assert_array_equal(result['var_level_10m'].values, [1, 3, 5])
     np.testing.assert_array_equal(result['var_level_20m'].values, [2, 4, 6])
 
-    # Create a sample dataset
-    data = np.array([[1, 2], [3, 4], [5, 6]])
-    ds = xr.Dataset({
-        'var': (('time', 'level'), data)
-    }, coords={'time': [0, 1, 2], 'level': [10, 20]})
-
     # Run the function with custom labels
     result = convert_2d_to_1d(ds, parse='level', dim_labels=['low', 'high'])
 
@@ -619,3 +594,16 @@ def test_convert_2d_to_1d():
     assert 'var_high' in result
     np.testing.assert_array_equal(result['var_low'].values, [1, 3, 5])
     np.testing.assert_array_equal(result['var_high'].values, [2, 4, 6])
+
+    # Create a sample dataset
+    data = np.array([[1], [3], [5]])
+    ds = xr.Dataset({
+        'var': (('time', 'level'), data)
+    }, coords={'time': [0, 1, 2], 'level': [10]})
+
+    # Run the function with keep_name_if_one=True
+    result = convert_2d_to_1d(ds, parse='level', keep_name_if_one=True)
+
+    # Check the results
+    assert 'var' in result
+    np.testing.assert_array_equal(result['var'].values, [1, 3, 5])