diff --git a/tests/unit/test_rain.py b/tests/unit/test_rain.py new file mode 100644 index 0000000..bf067b4 --- /dev/null +++ b/tests/unit/test_rain.py @@ -0,0 +1,37 @@ +import unittest +import sys +import os +import pandas as pd +import numpy as np + +# Add the src folder to the Python path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../src'))) + +from my_package.util import ensure_rain_column +from my_package.util import fill_rain_column + + +class TestRain(unittest.TestCase): + def test_ensure_rain(self): + # Creates a df with the column 'something' and the value NaN + self.df = pd.DataFrame([[np.nan]], columns=['something']) + + # Runs the function that should create a column 'rain.1h' if its not already existing + self.df = ensure_rain_column(self.df) + + # Checks if the column 'rain.1h' is in the columns of the dataframe + self.assertIn('rain.1h', self.df.columns) + + def test_fill_rain(self): + # Creates a dataframe with the column 'rain.1h' and value NaN + self.df = pd.DataFrame([[np.nan]], columns=['rain.1h']) + + # Runs the function, that should replace NaN with 0 + self.df = fill_rain_column(self.df) + + # Checks if the first index in the 'rain.1h' column is not equal to NaN, aka the value has changed + self.assertNotEqual(self.df['rain.1h'][0], np.nan) + +if __name__ == "__main__": + unittest.main() +