1 files changed, 39 insertions, 71 deletions
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