From c2c3a85b798cd8686d3d2c468d47123f5f3df469 Mon Sep 17 00:00:00 2001 From: Lopa10ko Date: Mon, 3 Feb 2025 15:42:12 +0300 Subject: [PATCH] fix: update finetune api test --- tests/unit/api/main/test_api_main.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/unit/api/main/test_api_main.py b/tests/unit/api/main/test_api_main.py index 1670bb9b5..5c97d74be 100644 --- a/tests/unit/api/main/test_api_main.py +++ b/tests/unit/api/main/test_api_main.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from fedot.core.pipelines.pipeline_builder import PipelineBuilder from matplotlib import get_backend, pyplot as plt from fedot_ind.api.main import FedotIndustrial @@ -131,18 +132,19 @@ def test_generate_anomaly_ts(ts_config, anomaly_config): assert interval[0] in ts_range and interval[1] in ts_range -@pytest.mark.parametrize('data_func, fedot_func', ( - [univariate_clf_data, fedot_industrial_classification], - [multivariate_clf_data, fedot_industrial_classification], - [univariate_regression_data, fedot_industrial_regression], - [multivariate_regression_data, fedot_industrial_regression], +@pytest.mark.parametrize('data_func, fedot_func, node', ( + [univariate_clf_data, fedot_industrial_classification, 'rf'], + [multivariate_clf_data, fedot_industrial_classification, 'rf'], + [univariate_regression_data, fedot_industrial_regression, 'treg'], + [multivariate_regression_data, fedot_industrial_regression, 'treg'], ), ids=['clf_uni', 'clf_multi', 'reg_uni', 'reg_multi']) -def test_finetune(data_func, fedot_func): +def test_finetune(data_func, fedot_func, node): data = data_func() fedot_industrial = fedot_func() - fedot_industrial.fit(data) - fedot_industrial.finetune(train_data=data, tuning_params={'tuning_timeout': 0.1}) - assert fedot_industrial.solver is not None + fedot_industrial.finetune(train_data=data, + model_to_tune=PipelineBuilder().add_node(node), + tuning_params={'tuning_timeout': 0.1}) + assert fedot_industrial.manager.solver is not None def test_plot_methods():