Add test for np.array in test_linear_extra_interpolation() and restructure

Author: SabineHaas

tests/test_tools.py

@@ -65,48 +65,68 @@ def test_smallest_difference(self):
     def test_linear_extra_interpolation(self):
         parameters = {'requested_height': 100,
                       'column_name': 'v_wind'}
-        data_frame = pd.DataFrame(data={'v_wind': [4.0, 5.0, 6.0]},
-                                  index=[100, 150, 200])
-        data_frame_series = pd.DataFrame(data={'v_wind': [
+        weather = pd.DataFrame(data={'v_wind': [4.0, 5.0, 6.0]},
+                               index=[100, 150, 200])
+        weather_series = pd.DataFrame(data={'v_wind': [
             pd.Series(data=[4.0, 5.0, 6.0]),
             pd.Series(data=[8.0, 10.0, 14.0]),
             pd.Series(data=[16.0, 20.0, 28.0])]}, index=[100, 150, 200])
-        # TODO: test v_wind as np.array
+        weather_arr = pd.DataFrame(data={'v_wind': [
+            np.array(weather_series['v_wind'][weather_series.index[0]]),
+            np.array(weather_series['v_wind'][weather_series.index[1]]),
+            np.array(weather_series['v_wind'][weather_series.index[2]])]},
+                                   index=weather_series.index)
 
         # requested_height is an index of data frame
         exp_output = 4.0
-        assert (linear_extra_interpolation(data_frame, **parameters) ==
+        assert (linear_extra_interpolation(weather, **parameters) ==
                 exp_output)
-        exp_output = pd.Series(data=[4.0, 5.0, 6.0])
-        assert_series_equal(linear_extra_interpolation(data_frame_series,
+        exp_series = pd.Series(data=[4.0, 5.0, 6.0])
+        assert_series_equal(linear_extra_interpolation(weather_series,
                                                        **parameters),
-                            exp_output)
+                            exp_series)
+        exp_arr = np.array(exp_series)
+        assert_array_equal(linear_extra_interpolation(weather_arr,
+                                                      **parameters),
+                           exp_arr)
         # requested_height is between indices of data frame
         exp_output = 5.5
         parameters['requested_height'] = 175
-        assert (linear_extra_interpolation(data_frame, **parameters) ==
+        assert (linear_extra_interpolation(weather, **parameters) ==
                 exp_output)
-        exp_output = pd.Series(data=[12.0, 15.0, 21.0])
-        assert_series_equal(linear_extra_interpolation(data_frame_series,
+        exp_series = pd.Series(data=[12.0, 15.0, 21.0])
+        assert_series_equal(linear_extra_interpolation(weather_series,
                                                        **parameters),
-                            exp_output)
+                            exp_series)
+        exp_arr = np.array(exp_series)
+        assert_array_equal(linear_extra_interpolation(weather_arr,
+                                                      **parameters),
+                           exp_arr)
         # requested_height > indices of data frame
         exp_output = 7.0
         parameters['requested_height'] = 250
-        assert (linear_extra_interpolation(data_frame, **parameters) ==
+        assert (linear_extra_interpolation(weather, **parameters) ==
                 exp_output)
-        exp_output = pd.Series(data=[24.0, 30.0, 42.0])
-        assert_series_equal(linear_extra_interpolation(data_frame_series,
+        exp_series = pd.Series(data=[24.0, 30.0, 42.0])
+        assert_series_equal(linear_extra_interpolation(weather_series,
                                                        **parameters),
-                            exp_output)
+                            exp_series)
+        exp_arr = np.array(exp_series)
+        assert_array_equal(linear_extra_interpolation(weather_arr,
+                                                      **parameters),
+                           exp_arr)
         # requested_height is < indices of data frame
         exp_output = 3.0
         parameters['requested_height'] = 50
-        assert (linear_extra_interpolation(data_frame, **parameters) ==
+        assert (linear_extra_interpolation(weather, **parameters) ==
                 exp_output)
         parameters['requested_height'] = 75
-        exp_output = pd.Series(data=[2.0, 2.5, 2.0])
-        assert_series_equal(linear_extra_interpolation(data_frame_series,
+        exp_series = pd.Series(data=[2.0, 2.5, 2.0])
+        assert_series_equal(linear_extra_interpolation(weather_series,
                                                        **parameters),
-                            exp_output)
+                            exp_series)
+        exp_arr = np.array(exp_series)
+        assert_array_equal(linear_extra_interpolation(weather_arr,
+                                                      **parameters),
+                           exp_arr)
         # TODO: Check this test