import unittest class MeanVectorAverageTests(unittest.TestCase): def _fut(self, input_windspd, input_winddir): from aosstower.wind import mean_wind_vector return mean_wind_vector(input_windspd, input_winddir) def test_scalar_input(self): for wdir, wspd in [(0, 20), (90, 20), (180, 20), (270, 20)]: output_speed, output_dir = self._fut(wspd, wdir) self.assertEqual(wdir, output_dir) self.assertEqual(wspd, 20) def test_array_input(self): for wdir, wspd in [([0, 90, 180, 270],[10, 10, 10, 10])]: output_speed, output_winddir = self._fut(wspd, wdir) self.assertEqual(wdir, output_winddir.tolist()) self.assertEqual(wspd, output_speed.tolist())