Modification test data

Author: Gsaes

qolmat/utils/data.py

@@ -11,6 +11,23 @@
 from qolmat.benchmark import missing_patterns
 
 
+def download_data(zipname: str, urllink: str, datapath: str = "data/") -> List[pd.DataFrame]:
+    path_zip = os.path.join(datapath)
+    print("Download data")
+    if not os.path.exists(path_zip + ".zip"):
+        if not os.path.exists(datapath):
+            os.mkdir(datapath)
+        request.urlretrieve(urllink + zipname + ".zip", path_zip + ".zip")
+
+    with zipfile.ZipFile(path_zip + ".zip", "r") as zip_ref:
+        zip_ref.extractall(path_zip)
+    data_folder = os.listdir(path_zip)
+    subfolder = os.path.join(path_zip, data_folder[0])
+    data_files = os.listdir(subfolder)
+    list_df = [pd.read_csv(os.path.join(subfolder, file)) for file in data_files]
+    return list_df
+
+
 def get_data(
     name_data: str = "Beijing", datapath: str = "data/", download: Optional[bool] = True
 ) -> pd.DataFrame:
@@ -32,19 +49,7 @@ def get_data(
     if name_data == "Beijing":
         urllink = "https://archive.ics.uci.edu/ml/machine-learning-databases/00501/"
         zipname = "PRSA2017_Data_20130301-20170228"
-        path_zip = os.path.join(datapath, zipname)
-
-        if not os.path.exists(path_zip + ".zip"):
-            if not os.path.exists(datapath):
-                os.mkdir(datapath)
-            request.urlretrieve(urllink + zipname + ".zip", path_zip + ".zip")
-
-        with zipfile.ZipFile(path_zip + ".zip", "r") as zip_ref:
-            zip_ref.extractall(path_zip)
-        data_folder = os.listdir(path_zip)
-        subfolder = os.path.join(path_zip, data_folder[0])
-        data_files = os.listdir(subfolder)
-        list_df = [pd.read_csv(os.path.join(subfolder, file)) for file in data_files]
+        list_df = download_data(zipname, urllink, datapath=datapath)
         list_df = [preprocess_data(df) for df in list_df]
         df = pd.concat(list_df)
         return df

tests/utils/test_data.py

@@ -5,6 +5,7 @@
 import pytest
 
 from qolmat.utils import data
+from pytest_mock.plugin import MockerFixture
 
 columns = ["No", "year", "month", "day", "hour", "a", "b", "wd", "station"]
 df = pd.DataFrame(
@@ -30,31 +31,25 @@
 
 
 @pytest.mark.parametrize("name_data", ["Beijing", "Artificial", "Bug"])
-def test_utils_data_get_data(name_data: str) -> None:
+def test_utils_data_get_data(name_data: str, mocker: MockerFixture) -> None:
+    mock_download = mocker.patch("qolmat.utils.data.download_data", return_value=[df])
+    mocker.patch("qolmat.utils.data.preprocess_data", return_value=df_preprocess)
+    try:
+        df_result = data.get_data(name_data=name_data)
+    except ValueError:
+        assert name_data not in ["Beijing", "Artificial"]
+        np.testing.assert_raises(ValueError, data.get_data, name_data)
+        return
+
     if name_data == "Beijing":
-        df = data.get_data(name_data=name_data)
-        expected_columns = [
-            "PM2.5",
-            "PM10",
-            "SO2",
-            "NO2",
-            "CO",
-            "O3",
-            "TEMP",
-            "PRES",
-            "DEWP",
-            "RAIN",
-            "WSPM",
-        ]
-        assert isinstance(df, pd.DataFrame)
-        assert df.columns.tolist() == expected_columns
+        assert mock_download.call_count == 1
+        pd.testing.assert_frame_equal(df_result, df_preprocess)
     elif name_data == "Artificial":
-        df = data.get_data(name_data=name_data)
         expected_columns = ["signal", "X", "A", "E"]
-        assert isinstance(df, pd.DataFrame)
-        assert df.columns.tolist() == expected_columns
+        assert isinstance(df_result, pd.DataFrame)
+        assert df_result.columns.tolist() == expected_columns
     else:
-        np.testing.assert_raises(ValueError, data.get_data, name_data)
+        assert False
 
 
 @pytest.mark.parametrize("df", [df])