updated unittests according to latest changes (#431)

oracle · Nov 16, 2023 · 8a780b5 · 8a780b5
2 parents 6899cbc + a86d887
commit 8a780b5
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 79 deletions.
diff --git a/tests/unitary/with_extras/operator/forecast/test_model_automlx.py b/tests/unitary/with_extras/operator/forecast/test_model_automlx.py
@@ -9,6 +9,7 @@
 import pandas as pd
 import datapane as dp
 from ads.opctl.operator.lowcode.forecast.model.automlx import AutoMLXOperatorModel
+from ads.opctl.operator.lowcode.forecast.model.forecast_datasets import ForecastDatasets
 
 from ads.opctl.operator.lowcode.forecast.operator_config import (
     ForecastOperatorConfig,
@@ -137,9 +138,9 @@ def test_automlx_for_unsorted_data(self, mock__call_pandas_fsspec):
             if filename == "primary.csv"
             else self.additional_data
         )
-        automlx = AutoMLXOperatorModel(self.config)
+        datasets = ForecastDatasets(self.config)
+        automlx = AutoMLXOperatorModel(self.config, datasets)
 
-        automlx._load_data()
         outputs = automlx._build_model()
         self.assertFalse(outputs.empty)
 

diff --git a/tests/unitary/with_extras/operator/forecast/test_model_autots.py b/tests/unitary/with_extras/operator/forecast/test_model_autots.py
@@ -11,6 +11,7 @@
 import datapane as dp
 import autots
 from ads.opctl.operator.common.utils import _build_image, _parse_input_args
+from ads.opctl.operator.lowcode.forecast.model.forecast_datasets import ForecastDatasets
 from ads.opctl.operator.lowcode.forecast.model.autots import (
     AutoTSOperatorModel,
     AUTOTS_MAX_GENERATION,
@@ -39,19 +40,26 @@ def setUp(self):
         spec.tuning = None
         spec.model_kwargs = {}
         spec.confidence_interval_width = 0.7
+        spec.additional_data = None
         self.spec = spec
 
         config = Mock(spec=ForecastOperatorConfig)
         config.spec = self.spec
         self.config = config
 
+        datasets = Mock(spec=ForecastDatasets)
+        datasets.original_user_data = None
+        datasets.original_total_data = None
+        datasets.original_additional_data = None
+        datasets.full_data_dict = {}
+        datasets.target_columns = []
+        datasets.categories = []
+        self.datasets = datasets
+
     @patch("autots.AutoTS")
     @patch("pandas.concat")
     def test_autots_parameter_passthrough(self, mock_concat, mock_autots):
-        autots = AutoTSOperatorModel(self.config)
-        autots.full_data_dict = {}
-        autots.target_columns = []
-        autots.categories = []
+        autots = AutoTSOperatorModel(self.config, self.datasets)
         autots._build_model()
 
         # When model_kwargs does not have anything, defaults should be sent as parameters.