fix tests

etna-team · Sep 4, 2024 · 000690e · 000690e
1 parent fee27c3
commit 000690e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git 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
@@ -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)