extra test

pola-rs · Apr 9, 2024 · fe61fc2 · fe61fc2
1 parent 81cd34a
commit fe61fc2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/py-polars/polars/functions/business.py b/py-polars/polars/functions/business.py
@@ -36,10 +36,10 @@ def _make_week_mask(
     if weekend is None:
         return tuple([True] * 7)
     if isinstance(weekend, str):
-        weekend_set = {weekend}
+        weekend = {weekend}
     else:
-        weekend_set = set(weekend)
-    for day in weekend_set:
+        weekend = set(weekend)
+    for day in weekend:
         if day not in _day_names():
             msg = f"Expected one of {_day_names()}, got: {day}"
             raise ValueError(msg)

diff --git a/py-polars/tests/unit/functions/business/test_business_day_count.py b/py-polars/tests/unit/functions/business/test_business_day_count.py
@@ -64,6 +64,13 @@ def test_business_day_count_w_weekend() -> None:
     )["business_day_count"]
     expected = pl.Series("business_day_count", [1, 7], pl.Int32)
     assert_series_equal(result, expected)
+
+    result = df.select(
+        business_day_count=pl.business_day_count("start", "end", weekend=("Sun",)),
+    )["business_day_count"]
+    expected = pl.Series("business_day_count", [1, 7], pl.Int32)
+    assert_series_equal(result, expected)
+
     result = df.select(
         business_day_count=pl.business_day_count(
             "start", "end", weekend=("Thu", "Fri", "Sat")