From cf239865ce440e1a71ee2360514eaeb018620ac9 Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Wed, 26 Aug 2026 16:46:45 -0700 Subject: Ruff afternoon! --- tests/models/spectrum/test_question.py | 5 ++ tests/models/spectrum/test_survey.py | 48 ++++++------- tests/models/spectrum/test_survey_manager.py | 103 +++++++++------------------ 3 files changed, 61 insertions(+), 95 deletions(-) (limited to 'tests/models/spectrum') diff --git a/tests/models/spectrum/test_question.py b/tests/models/spectrum/test_question.py index 57d260d..a44286d 100644 --- a/tests/models/spectrum/test_question.py +++ b/tests/models/spectrum/test_question.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import UTC, datetime from generalresearch.models import Source @@ -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", @@ -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", diff --git a/tests/models/spectrum/test_survey.py b/tests/models/spectrum/test_survey.py index 7365c7e..7ddd407 100644 --- a/tests/models/spectrum/test_survey.py +++ b/tests/models/spectrum/test_survey.py @@ -1,15 +1,25 @@ +from __future__ import annotations + from datetime import UTC, datetime from decimal import Decimal +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 + 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,18 +122,6 @@ 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 = { @@ -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"], @@ -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 @@ -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"]), diff --git a/tests/models/spectrum/test_survey_manager.py b/tests/models/spectrum/test_survey_manager.py index ce26c44..11dc01f 100644 --- a/tests/models/spectrum/test_survey_manager.py +++ b/tests/models/spectrum/test_survey_manager.py @@ -1,69 +1,32 @@ -import copy +from __future__ import annotations + import logging from datetime import UTC, datetime from decimal import Decimal +from typing import Any from pymysql import IntegrityError -logger = logging.getLogger() +from generalresearch.config import is_debug +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=UTC), - "category": "Exciting New", - "category_code": 232, - "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, - "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=UTC) spectrum_rw.execute_sql_query( @@ -73,24 +36,29 @@ 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=UTC) spectrum_rw.execute_sql_query( @@ -100,14 +68,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 @@ -122,8 +89,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 -- cgit v1.2.3