Add simple tests

Author: lesteve

setup.py

@@ -1,5 +1,7 @@
 import os
 from datetime import datetime, MAXYEAR
+from collections import namedtuple
+
 from setuptools import setup
 from setuptools.command.install import install
 
@@ -8,6 +10,33 @@
     LONG_DESCRIPTION = f.read()
 
 
+def get_brownout_schedule():
+    all_start_dates = [datetime(2022, 12, 1)] + [
+        datetime(2023, 2 * i, 1) for i in range(1, 7)
+    ]
+    date_in_the_far_future = datetime(MAXYEAR, 12, 31, 23, 59, 59)
+    all_end_dates = all_start_dates[1:] + [date_in_the_far_future]
+    all_check_functions = [
+        lambda dt: dt.minute < 5,
+        lambda dt: dt.minute < 10,
+        lambda dt: dt.minute < 15,
+        lambda dt: (dt.minute < 10) or (30 <= dt.minute < 40),
+        lambda dt: (dt.minute < 15) or (30 <= dt.minute < 45),
+        lambda dt: (dt.minute < 20) or (30 <= dt.minute < 50),
+        lambda dt: True,
+    ]
+
+    CheckedDatetimeWindow = namedtuple(
+        "CheckedDatetimeWindow", "start_datetime end_datetime check_function"
+    )
+    brownout_schedule = [
+        CheckedDatetimeWindow(*each)
+        for each in zip(all_start_dates, all_end_dates, all_check_functions)
+    ]
+
+    return brownout_schedule
+
+
 def maybe_raise_error(checked_datetime):
     accept_deprecated_sklearn_package_install = os.environ.get(
         "SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL", "use-brownout-schedule"
@@ -25,22 +54,6 @@ def maybe_raise_error(checked_datetime):
             "because SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=False is set"
         )
 
-    all_start_dates = [datetime(2022, 12, 1)] + [
-        datetime(2023, 2 * i, 1) for i in range(1, 7)
-    ]
-    date_in_the_far_future = datetime(MAXYEAR, 12, 31, 23, 59, 59)
-    all_end_dates = all_start_dates[1:] + [date_in_the_far_future]
-    all_check_functions = [
-        lambda dt: dt.minute < 5,
-        lambda dt: dt.minute < 10,
-        lambda dt: dt.minute < 15,
-        lambda dt: (dt.minute < 10) or (30 < dt.minute < 40),
-        lambda dt: (dt.minute < 15) or (30 < dt.minute < 45),
-        lambda dt: (dt.minute < 20) or (30 < dt.minute < 50),
-        lambda dt: True,
-    ]
-    brownout_schedule = zip(all_start_dates, all_end_dates, all_check_functions)
-
     error_message = "\n".join(
         [
             "The 'sklearn' PyPI package is deprecated, use 'scikit-learn' ",
@@ -63,6 +76,8 @@ def maybe_raise_error(checked_datetime):
             "https://github.com/scikit-learn/sklearn-pypi-package/issues/new",
         ]
     )
+
+    brownout_schedule = get_brownout_schedule()
     for start_date, end_date, check_function in brownout_schedule:
         if (start_date <= checked_datetime < end_date) and check_function(
             checked_datetime

test_brownout.py

@@ -0,0 +1,93 @@
+from datetime import timedelta
+from datetime import datetime
+
+import pytest
+
+from setup import get_brownout_schedule
+from setup import maybe_raise_error
+
+
+def test_brownout():
+    brownout_schedule = get_brownout_schedule()
+
+    brownout_iter = iter(brownout_schedule)
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=4))
+    maybe_raise_error(start_datetime + timedelta(minutes=5))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=9))
+
+    maybe_raise_error(start_datetime + timedelta(minutes=10))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=14))
+    maybe_raise_error(start_datetime + timedelta(minutes=15))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=9))
+    maybe_raise_error(start_datetime + timedelta(minutes=10))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=30))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=39))
+    maybe_raise_error(start_datetime + timedelta(minutes=40))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=14))
+    maybe_raise_error(start_datetime + timedelta(minutes=15))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=30))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=44))
+    maybe_raise_error(start_datetime + timedelta(minutes=45))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=19))
+    maybe_raise_error(start_datetime + timedelta(minutes=20))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=30))
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=49))
+    maybe_raise_error(start_datetime + timedelta(minutes=50))
+
+    start_datetime = next(brownout_iter).start_datetime
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(start_datetime + timedelta(minutes=59))
+
+    checked_datetime = datetime(2030, 1, 1)
+    with pytest.raises(SystemExit):
+        maybe_raise_error(checked_datetime)
+
+
+def test_allow_sklearn_package_install_environment_variable(monkeypatch):
+    with monkeypatch.context() as context:
+        context.setenv("SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL", "True")
+        checked_datetime = datetime(2030, 1, 1)
+        maybe_raise_error(checked_datetime)
+
+        context.setenv("SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL", "False")
+        checked_datetime = datetime(2020, 1, 1)
+        with pytest.raises(SystemExit):
+            maybe_raise_error(checked_datetime)