From 000690e6b3a457e80c08e38cbe7aa8741c809419 Mon Sep 17 00:00:00 2001 From: Vitalii Fuglaev Date: Wed, 4 Sep 2024 23:54:25 +0300 Subject: [PATCH] fix tests --- tests/test_commands/test_backtest.py | 12 ++++++------ tests/test_commands/test_forecast.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_commands/test_backtest.py b/tests/test_commands/test_backtest.py index d4530559d..c0a6669dd 100644 --- a/tests/test_commands/test_backtest.py +++ b/tests/test_commands/test_backtest.py @@ -188,7 +188,7 @@ def test_backtest_with_numeric_segments( base_timeseries_numeric_segments_path, ): target = pd.read_csv(base_timeseries_numeric_segments_path, dtype={"segment": str}) - segments = target.segments.unique() + segments = target["segment"].unique() tmp_output = TemporaryDirectory() tmp_output_path = Path(tmp_output.name) @@ -204,8 +204,8 @@ def test_backtest_with_numeric_segments( ] ) df_forecast = pd.read_csv(tmp_output_path / "forecast.csv", dtype={"segment": str}) - output_segments = df_forecast.segment.unique() - assert segments == output_segments + output_segments = df_forecast["segment"].unique() + assert set(segments) == set(output_segments) def test_backtest_with_numeric_segments_with_exog( @@ -215,7 +215,7 @@ def test_backtest_with_numeric_segments_with_exog( base_timeseries_numeric_segments_exog_path, ): target = pd.read_csv(base_timeseries_numeric_segments_path, dtype={"segment": str}) - segments = target.segments.unique() + segments = target["segment"].unique() tmp_output = TemporaryDirectory() tmp_output_path = Path(tmp_output.name) @@ -232,5 +232,5 @@ def test_backtest_with_numeric_segments_with_exog( ] ) df_forecast = pd.read_csv(tmp_output_path / "forecast.csv", dtype={"segment": str}) - output_segments = df_forecast.segment.unique() - assert segments == output_segments + output_segments = df_forecast["segment"].unique() + assert set(segments) == set(output_segments) diff --git a/tests/test_commands/test_forecast.py b/tests/test_commands/test_forecast.py index 5938d817c..be737635a 100644 --- a/tests/test_commands/test_forecast.py +++ b/tests/test_commands/test_forecast.py @@ -416,7 +416,7 @@ def test_forecast_with_numeric_segments( base_timeseries_numeric_segments_path, ): target = pd.read_csv(base_timeseries_numeric_segments_path, dtype={"segment": str}) - segments = target.segments.unique() + segments = target["segment"].unique() tmp_output = NamedTemporaryFile("w") tmp_output_path = Path(tmp_output.name) @@ -431,8 +431,8 @@ def test_forecast_with_numeric_segments( ], ) df_output = pd.read_csv(tmp_output_path, dtype={"segment": str}) - output_segments = df_output.segment.unique() - assert segments == output_segments + output_segments = df_output["segment"].unique() + assert set(segments) == set(output_segments) @pytest.mark.parametrize( @@ -446,7 +446,7 @@ def test_forecast_with_numeric_segments_with_exog( request, ): target = pd.read_csv(base_timeseries_numeric_segments_path, dtype={"segment": str}) - segments = target.segments.unique() + segments = target["segment"].unique() tmp_output = NamedTemporaryFile("w") tmp_output_path = Path(tmp_output.name) @@ -463,5 +463,5 @@ def test_forecast_with_numeric_segments_with_exog( ], ) df_output = pd.read_csv(tmp_output_path, dtype={"segment": str}) - output_segments = df_output.segment.unique() - assert segments == output_segments + output_segments = df_output["segment"].unique() + assert set(segments) == set(output_segments)