Merge branch 'master' into dev-greg

3 files changed, 98 insertions, 119 deletions

diff --git a/tests/models/spectrum/test_question.py b/tests/models/spectrum/test_question.py
index ba118d7..d469530 100644
--- a/tests/models/spectrum/test_question.py
+++ b/tests/models/spectrum/test_question.py
@@ -1,17 +1,19 @@
-from datetime import datetime, timezone
+from __future__ import annotations
 
-from generalresearch.models import Source
+from datetime import UTC, datetime
+
+from generalresearch.models.definitions import Source
 from generalresearch.models.spectrum.question import (
-    SpectrumQuestionOption,
     SpectrumQuestion,
-    SpectrumQuestionType,
     SpectrumQuestionClass,
+    SpectrumQuestionOption,
+    SpectrumQuestionType,
 )
 from generalresearch.models.thl.profiling.upk_question import (
     UpkQuestion,
+    UpkQuestionChoice,
     UpkQuestionSelectorMC,
     UpkQuestionType,
-    UpkQuestionChoice,
 )
 
 
@@ -32,6 +34,7 @@ class TestSpectrumQuestion:
             "mod_on": 1706557247467,
         }
         q = SpectrumQuestion.from_api(example_1, "us", "eng")
+        assert isinstance(q, SpectrumQuestion)
 
         expected_q = SpectrumQuestion(
             question_id="213",
@@ -43,7 +46,7 @@ class TestSpectrumQuestion:
             tags=None,
             options=None,
             class_num=SpectrumQuestionClass.CORE,
-            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc),
+            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC),
             is_live=True,
             source=Source.SPECTRUM,
             category_id=None,
@@ -72,6 +75,8 @@ class TestSpectrumQuestion:
             "mod_on": 1706557249817,
         }
         q = SpectrumQuestion.from_api(example_2, "us", "eng")
+        assert isinstance(q, SpectrumQuestion)
+
         expected_q = SpectrumQuestion(
             question_id="211",
             country_iso="us",
@@ -85,7 +90,7 @@ class TestSpectrumQuestion:
                 SpectrumQuestionOption(id="112", text="Female", order=1),
             ],
             class_num=SpectrumQuestionClass.CORE,
-            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc),
+            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC),
             is_live=True,
             source=Source.SPECTRUM,
             category_id=None,
@@ -160,7 +165,7 @@ class TestSpectrumQuestion:
                 SpectrumQuestionOption(id="999", text="None of the above", order=3),
             ],
             class_num=SpectrumQuestionClass.EXTENDED,
-            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc),
+            created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC),
             is_live=True,
             source=Source.SPECTRUM,
             category_id=None,
diff --git a/tests/models/spectrum/test_survey.py b/tests/models/spectrum/test_survey.py
index b612a63..02c5d3f 100644
--- a/tests/models/spectrum/test_survey.py
+++ b/tests/models/spectrum/test_survey.py
@@ -1,15 +1,25 @@
-from datetime import datetime, timezone
+from __future__ import annotations
+
+from datetime import UTC, datetime
 from decimal import Decimal
 
+from generalresearch.models.definitions import (
+    LogicalOperator,
+    Source,
+    TaskCalculationType,
+)
+from generalresearch.models.spectrum import SpectrumStatus
+from generalresearch.models.spectrum.survey import (
+    SpectrumCondition,
+    SpectrumQuota,
+    SpectrumSurvey,
+)
+from generalresearch.models.thl.survey.condition import ConditionValueType
+
 
 class TestSpectrumCondition:
 
     def test_condition_create(self):
-        from generalresearch.models import LogicalOperator
-        from generalresearch.models.spectrum.survey import (
-            SpectrumCondition,
-        )
-        from generalresearch.models.thl.survey.condition import ConditionValueType
 
         c = SpectrumCondition.from_api(
             {
@@ -64,10 +74,6 @@ class TestSpectrumCondition:
 class TestSpectrumQuota:
 
     def test_quota_create(self):
-        from generalresearch.models.spectrum.survey import (
-            SpectrumCondition,
-            SpectrumQuota,
-        )
 
         d = {
             "quota_id": "a846b545-4449-4d76-93a2-f8ebdf6e711e",
@@ -84,9 +90,6 @@ class TestSpectrumQuota:
         assert q.is_open
 
     def test_quota_passes(self):
-        from generalresearch.models.spectrum.survey import (
-            SpectrumQuota,
-        )
 
         q = SpectrumQuota(remaining_count=57, condition_hashes=["a"])
         assert q.passes({"a": True})
@@ -103,9 +106,6 @@ class TestSpectrumQuota:
         assert not q.passes({"a": True})
 
     def test_quota_passes_soft(self):
-        from generalresearch.models.spectrum.survey import (
-            SpectrumQuota,
-        )
 
         q = SpectrumQuota(remaining_count=57, condition_hashes=["a", "b", "c"])
         # Pass if we match all
@@ -122,29 +122,17 @@ class TestSpectrumQuota:
 
 class TestSpectrumSurvey:
     def test_survey_create(self):
-        from generalresearch.models import (
-            LogicalOperator,
-            Source,
-            TaskCalculationType,
-        )
-        from generalresearch.models.spectrum import SpectrumStatus
-        from generalresearch.models.spectrum.survey import (
-            SpectrumCondition,
-            SpectrumQuota,
-            SpectrumSurvey,
-        )
-        from generalresearch.models.thl.survey.condition import ConditionValueType
 
         # Note: d is the raw response after calling SpectrumAPI.preprocess_survey() on it!
         d = {
             "survey_id": 29333264,
             "survey_name": "Exciting New Survey #29333264",
             "survey_status": 22,
-            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc),
+            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC),
             "category": "Exciting New",
             "category_code": 232,
-            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc),
-            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc),
+            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC),
+            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC),
             "soft_launch": False,
             "click_balancing": 0,
             "price_type": 1,
@@ -202,6 +190,8 @@ class TestSpectrumSurvey:
             "exclusion_period": 0,
         }
         s = SpectrumSurvey.from_api(d)
+        assert isinstance(s, SpectrumSurvey)
+
         expected_survey = SpectrumSurvey(
             cpi=Decimal("1.20000"),
             country_isos=["fr"],
@@ -212,7 +202,7 @@ class TestSpectrumSurvey:
             survey_id="29333264",
             survey_name="Exciting New Survey #29333264",
             status=SpectrumStatus.LIVE,
-            field_end_date=datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc),
+            field_end_date=datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC),
             category_code="232",
             calculation_type=TaskCalculationType.COMPLETES,
             requires_pii=False,
@@ -240,8 +230,8 @@ class TestSpectrumSurvey:
                     values=["18-64"],
                 )
             },
-            created_api=datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc),
-            modified_api=datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc),
+            created_api=datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC),
+            modified_api=datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC),
             updated=None,
         )
         assert expected_survey.model_dump_json() == s.model_dump_json()
@@ -255,11 +245,11 @@ class TestSpectrumSurvey:
             "survey_id": 29333264,
             "survey_name": "#29333264",
             "survey_status": 22,
-            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc),
+            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC),
             "category": "Exciting New",
             "category_code": 232,
-            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc),
-            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc),
+            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC),
+            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC),
             "soft_launch": False,
             "click_balancing": 0,
             "price_type": 1,
@@ -303,6 +293,8 @@ class TestSpectrumSurvey:
             "exclusion_period": 0,
         }
         s = SpectrumSurvey.from_api(d)
+        assert isinstance(s, SpectrumSurvey)
+
         assert {"212", "1202", "214"} == s.used_question_ids
         assert s.is_live
         assert s.is_open
@@ -318,11 +310,11 @@ class TestSpectrumSurvey:
             "survey_id": 29333264,
             "survey_name": "#29333264",
             "survey_status": 22,
-            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc),
+            "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC),
             "category": "Exciting New",
             "category_code": 232,
-            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc),
-            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc),
+            "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC),
+            "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC),
             "soft_launch": False,
             "click_balancing": 0,
             "price_type": 1,
@@ -345,6 +337,8 @@ class TestSpectrumSurvey:
             "exclusion_period": 0,
         }
         s = SpectrumSurvey.from_api(d)
+        assert isinstance(s, SpectrumSurvey)
+
         s.qualifications = ["a", "b", "c"]
         s.quotas = [
             SpectrumQuota(remaining_count=10, condition_hashes=["a", "b"]),
@@ -411,3 +405,15 @@ class TestSpectrumSurvey:
         assert (None, {"c", "d"}) == s.determine_eligibility_soft(
             {"a": True, "b": True, "c": None, "d": None}
         )
+
+
+def test_spectrum_something(
+    spectrum_conditions: list[SpectrumCondition], spectrum_api_surveys_json: list[str]
+):
+
+    c1 = spectrum_conditions[0]
+    c3 = spectrum_conditions[2]
+
+    survey = SpectrumSurvey.model_validate_json(spectrum_api_surveys_json[0])
+    assert c1.criterion_hash in survey.qualifications
+    assert c3.criterion_hash in survey.qualifications
diff --git a/tests/models/spectrum/test_survey_manager.py b/tests/models/spectrum/test_survey_manager.py
index 582093c..0300956 100644
--- a/tests/models/spectrum/test_survey_manager.py
+++ b/tests/models/spectrum/test_survey_manager.py
@@ -1,72 +1,36 @@
-import copy
+from __future__ import annotations
+
 import logging
-from datetime import timezone, datetime
+from datetime import UTC, datetime
 from decimal import Decimal
+from typing import TYPE_CHECKING, Any
 
 from pymysql import IntegrityError
 
+from generalresearch.config import is_debug
 
-logger = logging.getLogger()
+if TYPE_CHECKING:
+    from generalresearch.managers.spectrum.survey import (
+        SpectrumSurveyManager,
+    )
+    from generalresearch.sql_helper import SqlHelper
 
-example_survey_api_response = {
-    "survey_id": 29333264,
-    "survey_name": "#29333264",
-    "survey_status": 22,
-    "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc),
-    "category": "Exciting New",
-    "category_code": 232,
-    "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc),
-    "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc),
-    "soft_launch": False,
-    "click_balancing": 0,
-    "price_type": 1,
-    "pii": False,
-    "buyer_message": "",
-    "buyer_id": 4726,
-    "incl_excl": 0,
-    "cpi": Decimal("1.20"),
-    "last_complete_date": None,
-    "project_last_complete_date": None,
-    "quotas": [
-        {
-            "quota_id": "c2bc961e-4f26-4223-b409-ebe9165cfdf5",
-            "quantities": {"currently_open": 491, "remaining": 495, "achieved": 0},
-            "criteria": [
-                {
-                    "qualification_code": 214,
-                    "range_sets": [{"units": 311, "to": 64, "from": 18}],
-                }
-            ],
-        }
-    ],
-    "qualifications": [
-        {
-            "range_sets": [{"units": 311, "to": 64, "from": 18}],
-            "qualification_code": 212,
-        },
-        {"condition_codes": ["111", "117", "112"], "qualification_code": 1202},
-    ],
-    "country_iso": "fr",
-    "language_iso": "fre",
-    "bid_ir": 0.4,
-    "bid_loi": 600,
-    "overall_ir": None,
-    "overall_loi": None,
-    "last_block_ir": None,
-    "last_block_loi": None,
-    "survey_exclusions": set(),
-    "exclusion_period": 0,
-}
+logger = logging.getLogger()
 
 
 class TestSpectrumSurvey:
 
-    def test_survey_create(self, settings, spectrum_manager, spectrum_rw):
+    def test_survey_create(
+        self,
+        spectrum_survey_manager: SpectrumSurveyManager,
+        spectrum_rw: SqlHelper,
+        spectrum_api_survey_json: dict[str, Any],
+    ):
         from generalresearch.models.spectrum.survey import SpectrumSurvey
 
-        assert settings.debug, "CRITICAL: Do not run this on production."
+        assert is_debug(), "CRITICAL: Do not run this on production."
 
-        now = datetime.now(tz=timezone.utc)
+        now = datetime.now(tz=UTC)
         spectrum_rw.execute_sql_query(
             query=f"""
                 DELETE FROM `{spectrum_rw.db}`.spectrum_survey
@@ -74,26 +38,31 @@ class TestSpectrumSurvey:
             commit=True,
         )
 
-        d = example_survey_api_response.copy()
-        s = SpectrumSurvey.from_api(d)
-        spectrum_manager.create(s)
+        s = SpectrumSurvey.from_api(spectrum_api_survey_json)
+        assert isinstance(s, SpectrumSurvey)
+        spectrum_survey_manager.create(s)
 
-        surveys = spectrum_manager.get_survey_library(updated_since=now)
+        surveys = spectrum_survey_manager.get_survey_library(updated_since=now)
         assert len(surveys) == 1
         assert "29333264" == surveys[0].survey_id
         assert s.is_unchanged(surveys[0])
 
         try:
-            spectrum_manager.create(s)
+            spectrum_survey_manager.create(s)
         except IntegrityError as e:
             print(e.args)
 
-    def test_survey_update(self, settings, spectrum_manager, spectrum_rw):
+    def test_survey_update(
+        self,
+        spectrum_survey_manager: SpectrumSurveyManager,
+        spectrum_rw: SqlHelper,
+        spectrum_api_survey_json: dict[str, Any],
+    ):
         from generalresearch.models.spectrum.survey import SpectrumSurvey
 
-        assert settings.debug, "CRITICAL: Do not run this on production."
+        assert is_debug(), "CRITICAL: Do not run this on production."
 
-        now = datetime.now(tz=timezone.utc)
+        now = datetime.now(tz=UTC)
         spectrum_rw.execute_sql_query(
             query=f"""
                 DELETE FROM `{spectrum_rw.db}`.spectrum_survey
@@ -101,14 +70,13 @@ class TestSpectrumSurvey:
             """,
             commit=True,
         )
-        d = copy.deepcopy(example_survey_api_response)
-        s = SpectrumSurvey.from_api(d)
-        print(s)
+        s = SpectrumSurvey.from_api(spectrum_api_survey_json)
+        assert isinstance(s, SpectrumSurvey)
 
-        spectrum_manager.create(s)
+        spectrum_survey_manager.create(s)
         s.cpi = Decimal("0.50")
-        spectrum_manager.update([s])
-        surveys = spectrum_manager.get_survey_library(updated_since=now)
+        spectrum_survey_manager.update([s])
+        surveys = spectrum_survey_manager.get_survey_library(updated_since=now)
         assert len(surveys) == 1
         assert "29333264" == surveys[0].survey_id
         assert Decimal("0.50") == surveys[0].cpi
@@ -123,8 +91,8 @@ class TestSpectrumSurvey:
         s.bid_loi = None
         s.overall_loi = 1000
         s.last_block_loi = 1000
-        spectrum_manager.update([s])
-        surveys = spectrum_manager.get_survey_library(updated_since=now)
+        spectrum_survey_manager.update([s])
+        surveys = spectrum_survey_manager.get_survey_library(updated_since=now)
         assert 600 == surveys[0].bid_loi
         assert 1000 == surveys[0].overall_loi
         assert 1000 == surveys[0].last_block_loi