Add tests

tests/test_pipeline/test_autoregressive_pipeline.py

@@ -30,7 +30,6 @@
 from tests.test_pipeline.utils import assert_pipeline_equals_loaded_original
 from tests.test_pipeline.utils import assert_pipeline_forecasts_given_ts
 from tests.test_pipeline.utils import assert_pipeline_forecasts_given_ts_with_prediction_intervals
-from tests.utils import to_be_fixed
 
 DEFAULT_METRICS = [MAE(mode=MetricAggregationMode.per_segment)]
 
@@ -347,7 +346,6 @@ def test_forecast_given_ts_with_prediction_interval(model, transforms, example_t
     assert_pipeline_forecasts_given_ts_with_prediction_intervals(pipeline=pipeline, ts=example_tsds, horizon=horizon)
 
 
-@to_be_fixed(NotImplementedError, "Adding target components is not currently implemented!")
 @pytest.mark.parametrize(
     "model_fixture",
     (
@@ -357,14 +355,20 @@ def test_forecast_given_ts_with_prediction_interval(model, transforms, example_t
         "prediction_interval_context_required_dummy_model",
     ),
 )
-def test_forecast_return_components(example_tsds, model_fixture, request):
+def test_forecast_return_components(
+    example_tsds, model_fixture, request, expected_component_a=10, expected_component_b=90
+):
     model = request.getfixturevalue(model_fixture)
     pipeline = AutoRegressivePipeline(model=model, horizon=10)
     pipeline.fit(example_tsds)
     forecast = pipeline.forecast(return_components=True)
+    assert sorted(forecast.target_components_names) == sorted(["target_component_a", "target_component_b"])
+
+    target_components_df = TSDataset.to_flatten(forecast.get_target_components())
+    assert (target_components_df["target_component_a"] == expected_component_a).all()
+    assert (target_components_df["target_component_b"] == expected_component_b).all()
 
 
-@to_be_fixed(NotImplementedError, "Adding target components is not currently implemented!")
 @pytest.mark.parametrize(
     "model_fixture",
     (
@@ -374,8 +378,15 @@ def test_forecast_return_components(example_tsds, model_fixture, request):
         "prediction_interval_context_required_dummy_model",
     ),
 )
-def test_predict_return_components(example_tsds, model_fixture, request):
+def test_predict_return_components(
+    example_tsds, model_fixture, request, expected_component_a=20, expected_component_b=180
+):
     model = request.getfixturevalue(model_fixture)
     pipeline = AutoRegressivePipeline(model=model, horizon=10)
     pipeline.fit(example_tsds)
     forecast = pipeline.predict(ts=example_tsds, return_components=True)
+    assert sorted(forecast.target_components_names) == sorted(["target_component_a", "target_component_b"])
+
+    target_components_df = TSDataset.to_flatten(forecast.get_target_components())
+    assert (target_components_df["target_component_a"] == expected_component_a).all()
+    assert (target_components_df["target_component_b"] == expected_component_b).all()