diff options
| -rw-r--r-- | generalresearch/managers/thl/survey.py | 14 | ||||
| -rw-r--r-- | generalresearch/models/custom_types.py | 1 | ||||
| -rw-r--r-- | generalresearch/models/innovate/survey.py | 19 | ||||
| -rw-r--r-- | generalresearch/models/thl/survey/model.py | 70 | ||||
| -rw-r--r-- | tests/managers/thl/test_survey.py | 33 | ||||
| -rw-r--r-- | tests/models/thl/test_survey.py | 140 |
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) + ] |