implement active hours on shared Task model

6 files changed, 269 insertions, 8 deletions

diff --git a/generalresearch/managers/thl/survey.py b/generalresearch/managers/thl/survey.py
index 92777e5..ada265f 100644
--- a/generalresearch/managers/thl/survey.py
+++ b/generalresearch/managers/thl/survey.py
@@ -91,10 +91,12 @@ class SurveyManager(PostgresManager):
         query = """
         INSERT INTO marketplace_survey (
             source, survey_id, created_at, updated_at,
-            is_live, is_recontact, buyer_id, eligibility_criteria
+            is_live, is_recontact, buyer_id, eligibility_criteria,
+            active_hours
         ) VALUES (
             %(source)s, %(survey_id)s, %(created_at)s, %(updated_at)s,
-            %(is_live)s, %(is_recontact)s, %(buyer_id)s, %(eligibility_criteria)s
+            %(is_live)s, %(is_recontact)s, %(buyer_id)s, %(eligibility_criteria)s,
+            %(active_hours)s
         ) ON CONFLICT (source, survey_id) DO NOTHING;"""
         params = [s.model_dump_sql() for s in surveys]
         with self.pg_config.make_connection() as conn:
@@ -216,7 +218,8 @@ class SurveyManager(PostgresManager):
     ) -> list[Survey]:
         query = """
         SELECT id, source, survey_id, created_at, updated_at,
-            is_live, is_recontact, buyer_id, eligibility_criteria
+            is_live, is_recontact, buyer_id, eligibility_criteria,
+            active_hours
         FROM marketplace_survey
         WHERE source = %(source)s AND
             survey_id = ANY(%(survey_ids)s);
@@ -233,7 +236,8 @@ class SurveyManager(PostgresManager):
         """
         query = """
         SELECT id, source, survey_id, created_at, updated_at,
-            is_live, is_recontact, buyer_id, eligibility_criteria
+            is_live, is_recontact, buyer_id, eligibility_criteria,
+            active_hours
         FROM marketplace_survey
         WHERE source = %(source)s AND is_live;
         """
@@ -245,7 +249,7 @@ class SurveyManager(PostgresManager):
         Return all live surveys
         """
         fields_default = """id, source, survey_id, created_at, updated_at,
-            is_live, is_recontact, buyer_id, eligibility_criteria"""
+            is_live, is_recontact, buyer_id, eligibility_criteria, active_hours"""
         fields = ", ".join(fields) if fields else fields_default
         query = f"""
         SELECT {fields}
diff --git a/generalresearch/models/custom_types.py b/generalresearch/models/custom_types.py
index 26c8b7e..5ac2978 100644
--- a/generalresearch/models/custom_types.py
+++ b/generalresearch/models/custom_types.py
@@ -202,6 +202,7 @@ DeviceTypes = Annotated[set[DeviceType], enum_to_comma_sep_str, from_comma_sep_s
 
 # This is a set of alphanumeric strings, that serializes and de-serializes
 # into a (sorted) comma-separated str
+# 2026-11-03 GSS: Nothing in here checks if they are alphanumeric... Should they be??
 AlphaNumStr = Annotated[str, StringConstraints(max_length=32, min_length=1)]
 
 # a string like an IP address, but we don't need to validate that it is
diff --git a/generalresearch/models/innovate/survey.py b/generalresearch/models/innovate/survey.py
index 3ca990c..03daea0 100644
--- a/generalresearch/models/innovate/survey.py
+++ b/generalresearch/models/innovate/survey.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 import logging
-from datetime import UTC, date
+from datetime import UTC, date, datetime
 from decimal import Decimal
 from functools import cached_property
 from typing import (
@@ -46,6 +46,7 @@ from generalresearch.models.thl.survey.condition import (
     ConditionValueType,
     MarketplaceCondition,
 )
+from generalresearch.models.thl.survey.model import TimeInterval
 
 logging.basicConfig()
 logger = logging.getLogger()
@@ -374,6 +375,22 @@ class InnovateSurvey(MarketplaceTask):
         d["off_hour_traffic"] = json.loads(d["off_hour_traffic"])
         return cls.model_validate(d)
 
+    def to_active_hours(self) -> list[TimeInterval] | None:
+        if self.off_hour_traffic is None:
+            return None
+        if "from" not in self.off_hour_traffic or "to" not in self.off_hour_traffic:
+            return None
+        # e.g. {"from": "01:00:00 AM", "to": "05:00:00 AM"}
+        start = datetime.strptime(self.off_hour_traffic["from"], "%I:%M:%S %p").time()  # noqa: DTZ007
+        end = datetime.strptime(self.off_hour_traffic["to"], "%I:%M:%S %p").time()  # noqa: DTZ007
+        return [
+            TimeInterval(
+                start=end,
+                end=start,
+                timezone=None,
+            )
+        ]
+
     def participation_allowed(
         self, att_survey_ids: set[str], att_job_ids: set[str]
     ) -> bool:
diff --git a/generalresearch/models/thl/survey/model.py b/generalresearch/models/thl/survey/model.py
index 9e3c03d..2ad2162 100644
--- a/generalresearch/models/thl/survey/model.py
+++ b/generalresearch/models/thl/survey/model.py
@@ -1,8 +1,10 @@
 from __future__ import annotations
 
-from datetime import UTC, datetime
+import json
+from datetime import UTC, datetime, time
 from decimal import Decimal
 from typing import Annotated, Any
+from zoneinfo import ZoneInfo
 
 from pydantic import (
     BaseModel,
@@ -15,6 +17,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
+from pydantic_extra_types.timezone_name import TimeZoneName
 
 from generalresearch.managers.thl.buyer import Buyer
 from generalresearch.models.custom_types import (
@@ -54,6 +57,48 @@ class SurveyEligibilityDefinition(BaseModel):
         return self
 
 
+class TimeInterval(BaseModel):
+    """
+    An optionally tz-aware time interval.
+    If timezone is null, the time interval is in the WORKER'S local time.
+
+    e.g. 07:00 to 22:00 America/New_York
+    """
+
+    start: time = Field(description="Start time", examples=[time(7)])
+    end: time = Field(description="End time", examples=[time(22)])
+    timezone: TimeZoneName | None = Field(
+        description="If timezone is null, the time interval is in the WORKER'S local time.",
+        examples=["America/New_York"],
+    )
+
+    def to_str(self):
+        return f"{self.start} - {self.end} tz={self.timezone}"
+
+    def is_currently_active(self, worker_local_time: datetime) -> bool:
+        assert worker_local_time.tzinfo is not None
+        if self.timezone is None:
+            current_time = worker_local_time.time()
+        else:
+            current_time = worker_local_time.astimezone(ZoneInfo(self.timezone)).time()
+
+        if self.start <= self.end:
+            return self.start <= current_time <= self.end
+
+        # An interval whose end precedes its start crosses midnight.
+        return current_time >= self.start or current_time <= self.end
+
+    def is_currently_active_by_worker_tz(self, worker_tz: ZoneInfo | None) -> bool:
+        if self.timezone is None:
+            if worker_tz is None:
+                return False
+            current_time = datetime.now(tz=worker_tz)
+        else:
+            current_time = datetime.now(tz=ZoneInfo(self.timezone))
+
+        return self.is_currently_active(current_time)
+
+
 class Survey(BaseModel):
     model_config = ConfigDict(validate_assignment=True)
 
@@ -80,6 +125,20 @@ class Survey(BaseModel):
 
     eligibility_criteria: SurveyEligibilityDefinition | None = Field(default=None)
 
+    # Task only allows entrances during ANY of the time intervals specified
+    active_hours: list[TimeInterval] | None = Field(default=None)
+
+    @staticmethod
+    def active_hours_from_off_hours(off_hours: TimeInterval) -> list[TimeInterval]:
+        """Convert one off-hours interval to its complementary active hours."""
+        return [
+            TimeInterval(
+                start=off_hours.end,
+                end=off_hours.start,
+                timezone=off_hours.timezone,
+            )
+        ]
+
     @property
     def natural_key(self) -> SurveyKey:
         return f"{self.source.value}:{self.survey_id}"
@@ -110,8 +169,17 @@ class Survey(BaseModel):
         d["eligibility_criteria"] = None
         if self.eligibility_criteria is not None:
             d["eligibility_criteria"] = self.eligibility_criteria.model_dump_json()
+        if self.active_hours:
+            d['active_hours'] = json.dumps(d['active_hours'])
         return d
 
+    def is_currently_active(self, worker_local_time: datetime):
+        return (
+            any(h.is_currently_active(worker_local_time) for h in self.active_hours)
+            if self.active_hours
+            else True
+        )
+
 
 class SurveyStat(BaseModel):
     model_config = ConfigDict(validate_assignment=True)
diff --git a/tests/managers/thl/test_survey.py b/tests/managers/thl/test_survey.py
index e114b70..928fe6c 100644
--- a/tests/managers/thl/test_survey.py
+++ b/tests/managers/thl/test_survey.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import uuid
 from collections.abc import Callable
-from datetime import UTC, datetime
+from datetime import UTC, datetime, time
 from decimal import Decimal
 from typing import TYPE_CHECKING
 
@@ -23,6 +23,7 @@ from generalresearch.models.thl.survey.model import (
     SurveyCategoryModel,
     SurveyEligibilityDefinition,
     SurveyStat,
+    TimeInterval,
 )
 
 if TYPE_CHECKING:
@@ -112,6 +113,36 @@ class TestSurvey:
         assert res2[0] == res[0]
         assert len(res2) == len(surveys2)
 
+    def test_active_hours_round_trip(
+        self,
+        delete_buyers_surveys: Callable[..., None],
+        survey_manager: SurveyManager,
+    ):
+        delete_buyers_surveys()
+        active_hours = [
+            TimeInterval(start=time(5), end=time(1), timezone=None),
+            TimeInterval(
+                start=time(8),
+                end=time(17),
+                timezone="America/New_York",
+            ),
+        ]
+        survey = Survey(
+            source=Source.TESTING,
+            survey_id="active-hours",
+            buyer_code="active-hours-buyer",
+            active_hours=active_hours,
+        )
+
+        survey_manager.create_or_update([survey])
+        result = survey_manager.filter_by_natural_key(
+            source=survey.source,
+            survey_ids=[survey.survey_id],
+        )
+
+        assert len(result) == 1
+        assert result[0].active_hours == active_hours
+
     def test_category(self, survey_manager: SurveyManager):
         survey1 = Survey(id=562289, survey_id="a", source=Source.TESTING)
         survey2 = Survey(id=562290, survey_id="a", source=Source.TESTING)
diff --git a/tests/models/thl/test_survey.py b/tests/models/thl/test_survey.py
new file mode 100644
index 0000000..b283e50
--- /dev/null
+++ b/tests/models/thl/test_survey.py
@@ -0,0 +1,140 @@
+from datetime import UTC, datetime, time
+from unittest.mock import patch
+from zoneinfo import ZoneInfo
+
+import pytest
+
+from generalresearch.models.definitions import Source
+from generalresearch.models.thl.survey.model import Survey, TimeInterval
+
+
+@pytest.mark.parametrize(
+    ("hour", "minute", "expected"),
+    [
+        (6, 59, False),
+        (7, 0, True),
+        (12, 0, True),
+        (22, 0, True),
+        (22, 1, False),
+    ],
+)
+def test_is_currently_active_in_worker_local_time(hour, minute, expected):
+    interval = TimeInterval(start=time(7), end=time(22), timezone=None)
+    worker_time = datetime(
+        2026, 1, 15, hour, minute, tzinfo=ZoneInfo("America/Chicago")
+    )
+
+    assert interval.is_currently_active(worker_time) is expected
+
+
+@pytest.mark.parametrize(
+    ("hour", "expected"),
+    [
+        (21, False),
+        (22, True),
+        (23, True),
+        (0, True),
+        (6, True),
+        (7, False),
+    ],
+)
+def test_is_currently_active_for_interval_crossing_midnight(hour, expected):
+    interval = TimeInterval(start=time(22), end=time(6), timezone=None)
+    worker_time = datetime(2026, 1, 15, hour, tzinfo=UTC)
+
+    assert interval.is_currently_active(worker_time) is expected
+
+
+def test_is_currently_active_converts_to_interval_timezone():
+    interval = TimeInterval(
+        start=time(7),
+        end=time(9),
+        timezone="America/New_York",
+    )
+    # 08:00 in New York on this date.
+    worker_time = datetime(
+        2026, 1, 15, 7, tzinfo=ZoneInfo("America/Chicago")
+    )
+
+    assert interval.is_currently_active(worker_time) is True
+
+
+def test_is_currently_active_requires_aware_worker_time():
+    interval = TimeInterval(start=time(7), end=time(22), timezone=None)
+
+    with pytest.raises(AssertionError):
+        interval.is_currently_active(datetime(2026, 1, 15, 12))
+
+
+def test_is_currently_active_requires_worker_timezone_for_local_interval():
+    interval = TimeInterval(start=time(7), end=time(22), timezone=None)
+
+    assert interval.is_currently_active_by_worker_tz(worker_tz=None) is False
+
+
+@pytest.mark.parametrize(
+    ("interval_timezone", "worker_timezone"),
+    [
+        (None, ZoneInfo("America/Chicago")),
+        ("America/New_York", None),
+    ],
+)
+def test_is_currently_active_supports_interval_crossing_midnight(
+    interval_timezone, worker_timezone
+):
+    interval = TimeInterval(
+        start=time(22),
+        end=time(6),
+        timezone=interval_timezone,
+    )
+    current_timezone = (
+        ZoneInfo(interval_timezone) if interval_timezone else worker_timezone
+    )
+    current_time = datetime(2026, 1, 15, 23, tzinfo=current_timezone)
+
+    with patch("generalresearch.models.thl.survey.model.datetime") as mock_datetime:
+        mock_datetime.now.return_value = current_time
+
+        assert interval.is_currently_active_by_worker_tz(worker_timezone) is True
+        mock_datetime.now.assert_called_once_with(tz=current_timezone)
+
+
+@pytest.mark.parametrize(
+    ("off_start", "off_end", "active_start", "active_end"),
+    [
+        (time(1), time(5), time(5), time(1)),
+        (time(22), time(6), time(6), time(22)),
+    ],
+)
+def test_from_off_hours_returns_complementary_active_hours(
+    off_start, off_end, active_start, active_end
+):
+    off_hours = TimeInterval(
+        start=off_start,
+        end=off_end,
+        timezone="America/New_York",
+    )
+
+    active_hours = Survey.active_hours_from_off_hours(off_hours)
+
+    assert active_hours == [
+        TimeInterval(
+            start=active_start,
+            end=active_end,
+            timezone="America/New_York",
+        )
+    ]
+
+
+def test_from_off_hours_result_can_be_assigned_to_survey():
+    off_hours = TimeInterval(start=time(1), end=time(5), timezone=None)
+
+    survey = Survey(
+        source=Source.TESTING,
+        survey_id="off-hours-test",
+        active_hours=Survey.active_hours_from_off_hours(off_hours),
+    )
+
+    assert survey.active_hours == [
+        TimeInterval(start=time(5), end=time(1), timezone=None)
+    ]