test: Fix tsfresh return type (again) (#240)

functime-org · Jun 11, 2024 · ad626c7 · ad626c7
1 parent ca1d354
commit ad626c7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py
@@ -4,9 +4,7 @@
 
 import polars as pl
 import pytest
-import statsmodels.api as sm
 from polars.testing import assert_frame_equal
-from statsmodels.tsa.stattools import acf as sm_acf
 
 from functime.cross_validation import train_test_split
 from functime.evaluation import (
@@ -53,6 +51,8 @@ def commodities_backtest(commodities_dataset):
 
 @pytest.mark.skip("Values do not align up with scipy")
 def test_acf(commodities_dataset):
+    from statsmodels.tsa.stattools import acf as sm_acf
+
     y_train, _ = commodities_dataset
     entity_col, _, target_col = y_train.columns
     # Result
@@ -106,6 +106,8 @@ def test_acf(commodities_dataset):
 
 @pytest.mark.skip("Values do not align up with scipy")
 def test_ljung_box(commodities_dataset):
+    import statsmodels.api as sm
+
     y_train, _ = commodities_dataset
     entity_col, _, target_col = y_train.columns
     # Result

diff --git a/tests/test_tsfresh.py b/tests/test_tsfresh.py
@@ -943,13 +943,13 @@ def test_longest_streak_below_mean(S, res):
         pl.DataFrame({"a": S}).select(
             longest_streak_below_mean(pl.col("a")).alias("lengths")
         ),
-        pl.DataFrame(pl.Series("lengths", res, dtype=pl.UInt64)),
+        pl.DataFrame(pl.Series("lengths", res, dtype=pl.Int32)),
     )
     assert_frame_equal(
         pl.LazyFrame({"a": S})
         .select(longest_streak_below_mean(pl.col("a")).alias("lengths"))
         .collect(),
-        pl.DataFrame(pl.Series("lengths", res, dtype=pl.UInt64)),
+        pl.DataFrame(pl.Series("lengths", res, dtype=pl.Int32)),
     )
 
 
@@ -969,13 +969,13 @@ def test_longest_streak_above_mean(S, res):
         pl.DataFrame({"a": S}).select(
             longest_streak_above_mean(pl.col("a")).alias("lengths")
         ),
-        pl.DataFrame(pl.Series("lengths", res, dtype=pl.UInt64)),
+        pl.DataFrame(pl.Series("lengths", res, dtype=pl.Int32)),
     )
     assert_frame_equal(
         pl.LazyFrame({"a": S})
         .select(longest_streak_above_mean(pl.col("a")).alias("lengths"))
         .collect(),
-        pl.DataFrame(pl.Series("lengths", res, dtype=pl.UInt64)),
+        pl.DataFrame(pl.Series("lengths", res, dtype=pl.Int32)),
     )