From a7c4022a90d06421f7bb4186eeeb566cffe9d86b Mon Sep 17 00:00:00 2001 From: Lopa10ko Date: Tue, 23 Jul 2024 13:00:07 +0300 Subject: [PATCH] chore: bring back cgru and cgru in default pipeline tests --- .../real_applications/test_heavy_models.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/test/integration/real_applications/test_heavy_models.py b/test/integration/real_applications/test_heavy_models.py index 356cbc3613..4b8559f26e 100644 --- a/test/integration/real_applications/test_heavy_models.py +++ b/test/integration/real_applications/test_heavy_models.py @@ -5,26 +5,15 @@ from fedot.core.pipelines.pipeline_builder import PipelineBuilder -@pytest.mark.skip(reason="Fails due to the https://github.com/aimclub/FEDOT/issues/1240") -def test_cgru_forecasting(): +@pytest.mark.parametrize('pipeline', ( + PipelineBuilder().add_node('lagged', params={'window_size': 200}).add_node('cgru').build(), + cgru_pipeline(), +), ids=str) +def test_cgru_forecasting(pipeline): horizon = 5 - window_size = 200 - train_data, test_data = get_ts_data('salaries', horizon) - - pipeline = PipelineBuilder().add_node('lagged', params={'window_size': window_size}).add_node('cgru').build() - pipeline.fit(train_data) - predicted = pipeline.predict(test_data).predict[0] - - assert len(predicted) == horizon - - -@pytest.mark.skip(reason="Fails due to the https://github.com/aimclub/FEDOT/issues/1240") -def test_cgru_in_pipeline(): - horizon = 5 - train_data, test_data = train_data, test_data = get_ts_data('salaries', horizon) - - pipeline = cgru_pipeline() + train_data, test_data, _ = get_ts_data('salaries', horizon) pipeline.fit(train_data) - predicted = pipeline.predict(test_data).predict[0] + predicted = pipeline.predict(test_data).predict + assert predicted is not None assert len(predicted) == horizon