diff --git a/.github/bump_version.py b/.github/bump_version.py index 779a82e38..05092fdad 100644 --- a/.github/bump_version.py +++ b/.github/bump_version.py @@ -42,10 +42,9 @@ def bump_version(version: str, bump: str) -> str: major, minor, patch = (int(x) for x in version.split(".")) if bump == "major": return f"{major + 1}.0.0" - elif bump == "minor": + if bump == "minor": return f"{major}.{minor + 1}.0" - else: - return f"{major}.{minor}.{patch + 1}" + return f"{major}.{minor}.{patch + 1}" def update_file(path: Path, old_version: str, new_version: str): diff --git a/.github/find-api-model-versions.py b/.github/find-api-model-versions.py index 1e5a262a4..ece882b19 100644 --- a/.github/find-api-model-versions.py +++ b/.github/find-api-model-versions.py @@ -5,6 +5,7 @@ from importlib.metadata import version as distribution_version from policyengine.bundle import get_current_bundle + from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, POLICYENGINE_CORE_VERSION, diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 7c2fb0cbf..9b8659899 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -86,6 +86,8 @@ jobs: run: bash .github/scripts/install_locked_ruff.sh - name: Format check with ruff run: ruff format --check . + - name: Lint with ruff + run: ruff check . typecheck-v2: name: Type-check API v2 # Full CI runs only on the lowest unmerged pull request in a stack. diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 8ca8510a3..97295ae71 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -29,6 +29,8 @@ jobs: run: bash .github/scripts/install_locked_ruff.sh - name: Format check with ruff run: ruff format --check . + - name: Lint with ruff + run: ruff check . quality-guards: name: Quality guards diff --git a/AGENTS.md b/AGENTS.md index f47b70413..cf573a889 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -118,11 +118,11 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. ### Code quality and types -- `api.quality.alembic-revisions-unused-imports`: Alembic revision files are exempt from unused-import checks (F401), because autogenerated revisions import types they may not use. (advisory: ruff per-file ignore of F401 for `migrations/*/versions/*.py`, not blocking) +- `api.quality.alembic-revisions-unused-imports`: Alembic revision files are exempt from ruff lint rules, because they must stay exactly as Alembic generates them. (advisory: ruff per-file ignore of ALL for `migrations/*/versions/*.py`, not blocking) - `api.quality.environment-variables-load`: Required environment variables and credentials load in a deployed-like environment. (blocking: CI job `test_env_vars`) - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. (blocking: mypy) - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format) -- `api.quality.ruff-lint`: Changed Python files pass ruff check with the rules selected in pyproject.toml (ruff's defaults, made explicit); agents run it before committing, but CI does not yet. (advisory: ruff check (E4, E7, E9, F), not blocking) +- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, UP, W)) - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`) ### Observability diff --git a/changelog.d/ruff-safe-autofixes.changed.md b/changelog.d/ruff-safe-autofixes.changed.md new file mode 100644 index 000000000..6d3b3b72b --- /dev/null +++ b/changelog.d/ruff-safe-autofixes.changed.md @@ -0,0 +1 @@ +Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding. diff --git a/dashboard/experiments/gpt4_api_user.py b/dashboard/experiments/gpt4_api_user.py index 37a2656d4..61c13bd0e 100644 --- a/dashboard/experiments/gpt4_api_user.py +++ b/dashboard/experiments/gpt4_api_user.py @@ -1,11 +1,12 @@ -import streamlit as st +import json +import os +import time + import openai -from openai.embeddings_utils import get_embedding, cosine_similarity import pandas as pd -import json import requests -import time -import os +import streamlit as st +from openai.embeddings_utils import cosine_similarity, get_embedding # Point OpenAI to the API key openai.api_key = os.environ["OPENAI_API_KEY"] @@ -58,20 +59,18 @@ def embed(prompt, engine="text-embedding-ada-002"): def write_prompt(question, relevant_parameters): - prompt = f"""Question: {question} + return f"""Question: {question} Here's some metadata about the parameters that are relevant to your question: Relevant parameters: {relevant_parameters} Task: Convert the question into a reform of the syntax below. If not specified, assume the reform starts on 2023-01-01 and ends on 2024-01-01. - + Reform syntax you should write in: {{ "parameter_name": {{ "start_date.end_date": value }} }} e.g. "basic rate at 25%" -> {{ "hmrc.income_tax.rates.uk[0].rate": {{ "2023-01-01.2024-01-01": 0.25 }} }} Reply only with valid JSON. """ - return prompt - reform_generation_prompt = write_prompt(question, top5) @@ -80,11 +79,10 @@ def write_prompt(question, relevant_parameters): @st.cache_data def get_gpt4_response(prompt): - response = openai.ChatCompletion.create( + return openai.ChatCompletion.create( model="gpt-4", messages=[{"role": "user", "content": prompt}], )["choices"][0]["message"]["content"] - return response response = get_gpt4_response(reform_generation_prompt) @@ -117,9 +115,7 @@ def get_policyengine_impact(reform): time.sleep(5) response = requests.get(url).json() - impact = response["result"] - - return impact + return response["result"] # Now, get GPT4 to answer the question with the impact JSON. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index aea9bf695..42149b814 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -8,11 +8,11 @@ Every rule that governs `policyengine-api` is listed here with how it is enforce | Tier | enforced | ratcheted | advisory | doc-only | Total | | --- | ---: | ---: | ---: | ---: | ---: | -| mechanical | 40 | 0 | 2 | 12 | 54 | +| mechanical | 41 | 0 | 1 | 12 | 54 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | -**Blocking:** 40 of 54 rules. +**Blocking:** 41 of 54 rules. ## Gaps @@ -33,8 +33,7 @@ Every rule that governs `policyengine-api` is listed here with how it is enforce ### Advisory (checked but not blocking) -- `api.quality.alembic-revisions-unused-imports`: Alembic revision files are exempt from unused-import checks (F401), because autogenerated revisions import types they may not use. -- `api.quality.ruff-lint`: Changed Python files pass ruff check with the rules selected in pyproject.toml (ruff's defaults, made explicit); agents run it before committing, but CI does not yet. +- `api.quality.alembic-revisions-unused-imports`: Alembic revision files are exempt from ruff lint rules, because they must stay exactly as Alembic generates them. ### Unproven (no proof that the checker rejects a bad example) @@ -265,17 +264,17 @@ flowchart LR | Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof | | --- | --- | --- | --- | --- | --- | --- | --- | -| `api.quality.alembic-revisions-unused-imports` | Alembic revision files are exempt from unused-import checks (F401), because autogenerated revisions import types they may not use. | mechanical | advisory | no | ruff per-file ignore of F401 for `migrations/*/versions/*.py` | — | — | +| `api.quality.alembic-revisions-unused-imports` | Alembic revision files are exempt from ruff lint rules, because they must stay exactly as Alembic generates them. | mechanical | advisory | no | ruff per-file ignore of ALL for `migrations/*/versions/*.py` | — | — | | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — | | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | -| `api.quality.ruff-lint` | Changed Python files pass ruff check with the rules selected in pyproject.toml (ruff's defaults, made explicit); agents run it before committing, but CI does not yet. | mechanical | advisory | no | ruff check (E4, E7, E9, F) | — | — | +| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, UP, W) | `pr.yml:lint` | — | | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | ```mermaid flowchart LR r_api_quality_alembic_revisions_unused_imports("api.quality.alembic-revisions-unused-imports (advisory)") - e_ruff_per_file_ignore_of_F401_for_migrations___versions___py["ruff per-file ignore of F401 for migrations/*/versions/*.py"] + e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py["ruff per-file ignore of ALL for migrations/*/versions/*.py"] local_only>"not run in CI"] r_api_quality_environment_variables_load("api.quality.environment-variables-load (enforced)") e_CI_job_test_env_vars["CI job test_env_vars"] @@ -286,21 +285,21 @@ flowchart LR r_api_quality_ruff_format("api.quality.ruff-format (enforced)") e_ruff_format["ruff format"] j_CI_pr_yml_lint["CI pr.yml:lint"] - r_api_quality_ruff_lint("api.quality.ruff-lint (advisory)") - e_ruff_check__E4__E7__E9__F_["ruff check (E4, E7, E9, F)"] + r_api_quality_ruff_lint("api.quality.ruff-lint (enforced)") + e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__UP__W_["ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, UP, W)"] r_api_quality_tests_pass("api.quality.tests-pass (enforced)") e_CI_job_test["CI job test"] j_CI_pr_yml_test["CI pr.yml:test"] - r_api_quality_alembic_revisions_unused_imports --> e_ruff_per_file_ignore_of_F401_for_migrations___versions___py - e_ruff_per_file_ignore_of_F401_for_migrations___versions___py -.-> local_only + r_api_quality_alembic_revisions_unused_imports --> e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py + e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py -.-> local_only r_api_quality_environment_variables_load --> e_CI_job_test_env_vars e_CI_job_test_env_vars --> j_CI_pr_yml_test_env_vars r_api_quality_mypy_v2 --> e_mypy e_mypy --> j_CI_pr_yml_typecheck_v2 r_api_quality_ruff_format --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint - r_api_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F_ - e_ruff_check__E4__E7__E9__F_ -.-> local_only + r_api_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__UP__W_ + e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__UP__W_ --> j_CI_pr_yml_lint r_api_quality_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test ``` diff --git a/migrations/v1/env.py b/migrations/v1/env.py index 8fce659c9..1669dc5ed 100644 --- a/migrations/v1/env.py +++ b/migrations/v1/env.py @@ -1,7 +1,7 @@ """Alembic environment for the API v1 schema.""" -from logging.config import fileConfig import os +from logging.config import fileConfig from alembic import context from sqlalchemy import engine_from_config, pool @@ -9,7 +9,6 @@ from policyengine_api.data.v1_models import V1Base - config = context.config if config.config_file_name is not None: fileConfig(config.config_file_name) diff --git a/migrations/v2/env.py b/migrations/v2/env.py index 7a65ce625..1ab23e003 100644 --- a/migrations/v2/env.py +++ b/migrations/v2/env.py @@ -13,7 +13,6 @@ ) from policyengine_api.data.v2.models import V2_METADATA - config = context.config if config.config_file_name is not None: fileConfig(config.config_file_name) diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 2c0d23a35..881f104df 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -3,8 +3,8 @@ This is the main Flask app for the PolicyEngine API. """ -import time import sys +import time start_time = time.time() @@ -16,18 +16,19 @@ def log_timing(message): log_timing("Basic imports completed") -from flask_cors import CORS import flask +from flask_cors import CORS log_timing("Flask imports completed") +from policyengine_observability import instrument_flask + from policyengine_api.extensions import cache from policyengine_api.migration_logging import register_migration_request_logging +from policyengine_api.observability import runtime as observability_runtime from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.request_context import REQUEST_ID_HEADER -from policyengine_api.observability import runtime as observability_runtime from policyengine_api.runtime_cache.settings import load_runtime_cache_settings -from policyengine_observability import instrument_flask log_timing("Caching utilities import completed") @@ -55,9 +56,9 @@ def log_timing(message): from policyengine_api.routes.user_profile_routes import user_profile_bp log_timing("User profile routes import completed") -from policyengine_api.routes.simulation_routes import simulation_bp -from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.reform_impact_routes import reform_impact_bp +from policyengine_api.routes.report_output_routes import report_output_bp +from policyengine_api.routes.simulation_routes import simulation_bp from policyengine_api.routes.system_routes import system_bp log_timing("Initialising API...") diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index 408cd8aa5..9fc4fd2a0 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -2,18 +2,28 @@ from __future__ import annotations +import time from collections.abc import Callable from contextlib import asynccontextmanager -import time from a2wsgi import WSGIMiddleware from fastapi import FastAPI, Request from fastapi.exception_handlers import request_validation_exception_handler from fastapi.exceptions import RequestValidationError +from policyengine_observability import ObservabilityRuntime +from starlette.datastructures import MutableHeaders +from starlette.middleware.cors import CORSMiddleware +from starlette.middleware.gzip import GZipMiddleware +from starlette.responses import PlainTextResponse, Response +from starlette.routing import Match, Mount +from starlette.types import ASGIApp + from policyengine_api.constants import VERSION from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.fastapi_routes.health import build_core_health_router -from policyengine_api.fastapi_routes.health import build_readiness_router +from policyengine_api.fastapi_routes.health import ( + build_core_health_router, + build_readiness_router, +) from policyengine_api.fastapi_routes.metadata import build_metadata_router from policyengine_api.fastapi_routes.specification import ( build_specification_router, @@ -29,6 +39,10 @@ ) from policyengine_api.migration_logging import log_migration_request from policyengine_api.observability import get_runtime +from policyengine_api.observability.identifiers import ( + OBSERVABILITY_ID_HEADER, + normalize_observability_id, +) from policyengine_api.request_context import ( REQUEST_ID_HEADER, _asgi_incoming_observability_id, @@ -37,17 +51,6 @@ current_observability_id, generate_request_id, ) -from policyengine_observability import ObservabilityRuntime -from policyengine_api.observability.identifiers import ( - OBSERVABILITY_ID_HEADER, - normalize_observability_id, -) -from starlette.datastructures import MutableHeaders -from starlette.middleware.cors import CORSMiddleware -from starlette.middleware.gzip import GZipMiddleware -from starlette.responses import PlainTextResponse, Response -from starlette.routing import Match, Mount -from starlette.types import ASGIApp def _apply_request_id_header( diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 7af395ffc..c1f754edf 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -3,30 +3,30 @@ import json import logging from dataclasses import dataclass -from policyengine_core.taxbenefitsystems import TaxBenefitSystem -from typing import Any, Union -from policyengine_api.utils import get_safe_json +from datetime import date, datetime +from importlib.metadata import version as get_package_version +from pathlib import Path +from typing import Any + +import dpath +import pandas as pd +from policyengine_core.model_api import Enum, Reform from policyengine_core.parameters import ( - ParameterNode, Parameter, + ParameterNode, ParameterScale, ParameterScaleBracket, + get_parameter, ) -from policyengine_core.parameters import get_parameter -from importlib.metadata import version as get_package_version -from policyengine_core.model_api import Reform, Enum from policyengine_core.periods import instant -import dpath -import pandas as pd -from datetime import date, datetime -from pathlib import Path -from policyengine_api.data.congressional_districts import ( - build_congressional_district_metadata, -) +from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_api.constants import ( get_bundle_default_dataset_option, ) +from policyengine_api.data.congressional_districts import ( + build_congressional_district_metadata, +) from policyengine_api.services.household_calculation_service import CalculationResult from policyengine_api.spm import ( SPM_INPUT_ERROR_CODES, @@ -35,7 +35,7 @@ spm_error_detail, spm_metadata, ) - +from policyengine_api.utils import get_safe_json logger = logging.getLogger(__name__) @@ -441,7 +441,7 @@ def build_entities(self) -> dict: def calculate( self, household: dict, - reform: Union[dict, None], + reform: dict | None, spm: dict | None = None, spm_requested: bool = False, prepared: PreparedCountryCalculation | None = None, @@ -539,7 +539,7 @@ def calculate( def prepare_calculation( self, household: dict, - reform: Union[dict, None], + reform: dict | None, spm: dict | None = None, spm_requested: bool = False, ) -> PreparedCountryCalculation: @@ -565,7 +565,7 @@ def prepare_calculation( def _create_simulation( self, household: dict, - reform: Union[dict, None], + reform: dict | None, spm: dict | None = None, ): selection = normalize_spm_selection(getattr(self, "country_id", ""), spm) diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index 6dbac4be1..c78752304 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -741,11 +741,7 @@ def normalize_us_region(region: str) -> str: normalizes the format. Use _validate_us_region for validation. """ # Already has a valid prefix - return as-is - if ( - region.startswith("state/") - or region.startswith("place/") - or region.startswith("congressional_district/") - ): + if region.startswith(("state/", "place/", "congressional_district/")): return region # National level - no prefix needed diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 51cc0fcdf..79f83a173 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -10,13 +10,13 @@ from typing import Any from sqlalchemy import ( - BigInteger, CHAR, + JSON, + BigInteger, CheckConstraint, DateTime, Index, Integer, - JSON, SmallInteger, String, Text, diff --git a/policyengine_api/data/v2/catalog/catalog_selection.py b/policyengine_api/data/v2/catalog/catalog_selection.py index e4a32fb3c..f998cf224 100644 --- a/policyengine_api/data/v2/catalog/catalog_selection.py +++ b/policyengine_api/data/v2/catalog/catalog_selection.py @@ -13,7 +13,6 @@ TaxBenefitModelVersion, ) - SUPPORTED_V2_COUNTRY_IDS = frozenset({"us", "uk"}) diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index c8bedeeed..0f04a4c4c 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -2,14 +2,14 @@ from __future__ import annotations +import json +import math +import re from collections import Counter from collections.abc import Callable, Mapping, Sequence -from datetime import date, datetime, timedelta, timezone +from datetime import UTC, date, datetime, timedelta from enum import Enum from importlib import metadata as importlib_metadata -import json -import math -import re from typing import Any from urllib.parse import urlsplit from uuid import NAMESPACE_URL, UUID, uuid5 @@ -28,7 +28,6 @@ VariableRecord, ) - SUPPORTED_COUNTRIES = ("us", "uk") REQUIRED_DEPENDENCIES = ( "policyengine-core", @@ -179,8 +178,8 @@ def _aware_datetime(value: object, *, field_name: str) -> datetime: if not isinstance(value, datetime): raise CatalogExtractionError(f"{field_name} must be a date or datetime") if value.tzinfo is None: - return value.replace(tzinfo=timezone.utc) - return value.astimezone(timezone.utc) + return value.replace(tzinfo=UTC) + return value.astimezone(UTC) def _optional_aware_datetime(value: object, *, field_name: str) -> datetime | None: diff --git a/policyengine_api/data/v2/catalog/initialization.py b/policyengine_api/data/v2/catalog/initialization.py index 873ed6a28..94978a2e8 100644 --- a/policyengine_api/data/v2/catalog/initialization.py +++ b/policyengine_api/data/v2/catalog/initialization.py @@ -2,9 +2,9 @@ from __future__ import annotations -from collections.abc import Callable, Mapping import json import sys +from collections.abc import Callable, Mapping from sqlalchemy import Engine, create_engine from sqlalchemy.pool import NullPool diff --git a/policyengine_api/data/v2/catalog/publication.py b/policyengine_api/data/v2/catalog/publication.py index 443aa61a1..6155dd836 100644 --- a/policyengine_api/data/v2/catalog/publication.py +++ b/policyengine_api/data/v2/catalog/publication.py @@ -2,9 +2,9 @@ from __future__ import annotations -from collections.abc import Callable import logging import time +from collections.abc import Callable import sqlalchemy as sa from sqlalchemy import Connection, Engine @@ -32,7 +32,6 @@ Simulation, ) - # Stable application-defined PostgreSQL lock ID shared by all v2 catalog publishers. PUBLICATION_ADVISORY_LOCK_KEY = 8_629_020_026_090_001 diff --git a/policyengine_api/data/v2/catalog/publication_reconciliation.py b/policyengine_api/data/v2/catalog/publication_reconciliation.py index 8ebac6d6e..b4cb49f33 100644 --- a/policyengine_api/data/v2/catalog/publication_reconciliation.py +++ b/policyengine_api/data/v2/catalog/publication_reconciliation.py @@ -4,7 +4,8 @@ import sqlalchemy as sa from sqlalchemy import Connection, Select -from sqlalchemy.dialects.postgresql import JSONB, insert as postgresql_insert +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import insert as postgresql_insert from policyengine_api.data.v2.catalog.publication_staging import ( STAGE_CATALOG_DATASETS, @@ -30,7 +31,6 @@ Variable, ) - DATASETS = Dataset.__table__ MODELS = TaxBenefitModel.__table__ MODEL_VERSIONS = TaxBenefitModelVersion.__table__ diff --git a/policyengine_api/data/v2/catalog/publication_staging.py b/policyengine_api/data/v2/catalog/publication_staging.py index 325d3ac3c..b00f18ec4 100644 --- a/policyengine_api/data/v2/catalog/publication_staging.py +++ b/policyengine_api/data/v2/catalog/publication_staging.py @@ -4,9 +4,9 @@ from collections.abc import Callable, Iterable, Iterator, Sequence +import sqlalchemy as sa from psycopg import sql from psycopg.types.json import Jsonb -import sqlalchemy as sa from sqlalchemy import Connection, MetaData, Table from sqlalchemy.dialects.postgresql import JSONB @@ -19,7 +19,6 @@ iter_batches, ) - COPY_BATCH_SIZE = 10_000 STAGING_METADATA = MetaData() diff --git a/policyengine_api/data/v2/database.py b/policyengine_api/data/v2/database.py index 97215a526..ea14cd325 100644 --- a/policyengine_api/data/v2/database.py +++ b/policyengine_api/data/v2/database.py @@ -16,7 +16,6 @@ load_v2_runtime_database_settings, ) - DATABASE_POOL_RECYCLE_SECONDS = 1800 DATABASE_POOL_SIZE = 5 DATABASE_POOL_MAX_OVERFLOW = 5 diff --git a/policyengine_api/data/v2/household_migration_qualification.py b/policyengine_api/data/v2/household_migration_qualification.py index be29ac6dd..f017ece60 100644 --- a/policyengine_api/data/v2/household_migration_qualification.py +++ b/policyengine_api/data/v2/household_migration_qualification.py @@ -2,10 +2,10 @@ from __future__ import annotations -from collections.abc import Callable, Collection, Mapping -from dataclasses import dataclass import json import sys +from collections.abc import Callable, Collection, Mapping +from dataclasses import dataclass from typing import Protocol from alembic.config import Config @@ -29,7 +29,6 @@ load_v2_migration_database_settings, ) - HOUSEHOLD_MIGRATION_REVISION = "724b1b11a33e" V2_ALEMBIC_CONFIG = "alembic-v2.ini" diff --git a/policyengine_api/data/v2/metadata_validation.py b/policyengine_api/data/v2/metadata_validation.py index 20b2efa0b..a37e92aaf 100644 --- a/policyengine_api/data/v2/metadata_validation.py +++ b/policyengine_api/data/v2/metadata_validation.py @@ -2,7 +2,6 @@ from collections.abc import Iterable - # These names belong only to the existing v1 database domain. Names that are # intentionally shared between v1 and v2, such as `simulations` and # `user_policies`, are excluded from this targeted rejection set. diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 4bddc292e..d832cc879 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -2,10 +2,10 @@ from __future__ import annotations +import os from collections.abc import Mapping from dataclasses import dataclass from datetime import date -import os from sqlalchemy import Connection, MetaData, inspect, text from sqlalchemy.engine import URL, make_url diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index 1b512f9e2..4e408a7af 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -6,7 +6,6 @@ validate_v2_metadata_table_names, ) - # Apply deterministic names before any v2 table is declared. This is a local # SQLAlchemy metadata facility that SQLModel does not wrap directly. SQLModel.metadata.naming_convention = { @@ -18,6 +17,28 @@ } # Import in dependency order so every reviewed table registers exactly once. +from policyengine_api.data.v2.models.associations import ( # noqa: E402 + UserHouseholdAssociation, + UserPolicy, + UserReportAssociation, + UserSimulationAssociation, +) +from policyengine_api.data.v2.models.comparison_runs import ( # noqa: E402 + Stage12AggregationStatus, + Stage12ComparisonReport, + Stage12ComparisonSimulation, + Stage12ResultComparisonStatus, + Stage12RunStatus, + Stage12SimulationRole, +) +from policyengine_api.data.v2.models.household_mappings import ( # noqa: E402 + LegacyHouseholdMapping, +) +from policyengine_api.data.v2.models.households import ( # noqa: E402 + Household, + HouseholdJob, + HouseholdJobStatus, +) from policyengine_api.data.v2.models.metadata import ( # noqa: E402 Dataset, DatasetVersion, @@ -30,38 +51,10 @@ TaxBenefitModelVersion, Variable, ) -from policyengine_api.data.v2.models.users import User # noqa: E402 from policyengine_api.data.v2.models.policies import ( # noqa: E402 Dynamic, Policy, ) -from policyengine_api.data.v2.models.households import ( # noqa: E402 - Household, - HouseholdJob, - HouseholdJobStatus, -) -from policyengine_api.data.v2.models.household_mappings import ( # noqa: E402 - LegacyHouseholdMapping, -) -from policyengine_api.data.v2.models.simulations import ( # noqa: E402 - Simulation, - SimulationStatus, - SimulationType, -) -from policyengine_api.data.v2.models.comparison_runs import ( # noqa: E402 - Stage12AggregationStatus, - Stage12ComparisonReport, - Stage12ComparisonSimulation, - Stage12ResultComparisonStatus, - Stage12RunStatus, - Stage12SimulationRole, -) -from policyengine_api.data.v2.models.associations import ( # noqa: E402 - UserHouseholdAssociation, - UserPolicy, - UserReportAssociation, - UserSimulationAssociation, -) from policyengine_api.data.v2.models.policy_mappings import ( # noqa: E402 LegacyPolicyMapping, LegacyUserMapping, @@ -87,7 +80,12 @@ ReportRunStatus, ReportRunTrigger, ) - +from policyengine_api.data.v2.models.simulations import ( # noqa: E402 + Simulation, + SimulationStatus, + SimulationType, +) +from policyengine_api.data.v2.models.users import User # noqa: E402 V2_METADATA = SQLModel.metadata validate_v2_metadata_table_names(V2_METADATA.tables) diff --git a/policyengine_api/data/v2/models/base.py b/policyengine_api/data/v2/models/base.py index e612eff8b..23c3809a9 100644 --- a/policyengine_api/data/v2/models/base.py +++ b/policyengine_api/data/v2/models/base.py @@ -2,14 +2,13 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from types import MappingProxyType from uuid import UUID, uuid4 import sqlalchemy as sa from sqlmodel import Field, SQLModel - # These are the complete direct-SQLAlchemy categories permitted in the v2 # table layer. Each remains underneath canonical SQLModel table classes and # Field/Relationship declarations; no parallel declarative model exists. @@ -47,7 +46,7 @@ def utc_now() -> datetime: """Return an aware UTC timestamp for application-created rows.""" - return datetime.now(timezone.utc) + return datetime.now(UTC) class IdentifiedModel(SQLModel): diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 76ebb50b1..a6831d659 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -1,7 +1,7 @@ """Canonical SQLModel tables for v2 households and household jobs.""" from datetime import datetime -from enum import Enum +from enum import StrEnum from typing import TYPE_CHECKING, Any from uuid import UUID @@ -24,7 +24,7 @@ from policyengine_api.data.v2.models.simulations import Simulation -class HouseholdJobStatus(str, Enum): +class HouseholdJobStatus(StrEnum): PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index d1a797187..1b7b9edb7 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -1,7 +1,7 @@ """Canonical SQLModel tables for v2 model metadata, regions, and datasets.""" from datetime import datetime -from enum import Enum +from enum import StrEnum from typing import TYPE_CHECKING, Any, Optional from uuid import UUID @@ -20,7 +20,7 @@ from policyengine_api.data.v2.models.simulations import Simulation -class RegionType(str, Enum): +class RegionType(StrEnum): NATIONAL = "national" COUNTRY = "country" STATE = "state" diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 59ac9c4f6..1213ef907 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -1,19 +1,19 @@ """Canonical SQLModel tables for stable reports, runs, and run outputs.""" from datetime import datetime -from enum import Enum +from enum import StrEnum from typing import Any from uuid import UUID import sqlalchemy as sa from sqlmodel import Field, Relationship +from policyengine_api.data.v2.models.associations import UserReportAssociation from policyengine_api.data.v2.models.base import ( IdentifiedModel, TimestampedModel, enum_type, ) -from policyengine_api.data.v2.models.associations import UserReportAssociation from policyengine_api.data.v2.models.households import Household from policyengine_api.data.v2.models.metadata import ( Dataset, @@ -25,33 +25,33 @@ from policyengine_api.data.v2.models.users import User -class ReportRunStatus(str, Enum): +class ReportRunStatus(StrEnum): PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" FAILED = "failed" -class ReportRunTrigger(str, Enum): +class ReportRunTrigger(StrEnum): INITIAL = "initial" MANUAL = "manual" SYSTEM = "system" -class OutputStatus(str, Enum): +class OutputStatus(StrEnum): PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" FAILED = "failed" -class AggregateType(str, Enum): +class AggregateType(StrEnum): SUM = "sum" MEAN = "mean" COUNT = "count" -class DecileType(str, Enum): +class DecileType(StrEnum): INCOME = "income" WEALTH = "wealth" diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index d076ea4df..816d4896f 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -1,7 +1,7 @@ """Canonical SQLModel table for v2 simulations.""" from datetime import datetime -from enum import Enum +from enum import StrEnum from typing import TYPE_CHECKING, Any from uuid import UUID @@ -22,14 +22,14 @@ from policyengine_api.data.v2.models.reports import Report -class SimulationStatus(str, Enum): +class SimulationStatus(StrEnum): PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" FAILED = "failed" -class SimulationType(str, Enum): +class SimulationType(StrEnum): HOUSEHOLD = "household" ECONOMY = "economy" diff --git a/policyengine_api/data/v2/policy_migration_qualification.py b/policyengine_api/data/v2/policy_migration_qualification.py index e2224844e..91832823e 100644 --- a/policyengine_api/data/v2/policy_migration_qualification.py +++ b/policyengine_api/data/v2/policy_migration_qualification.py @@ -2,10 +2,10 @@ from __future__ import annotations -from collections.abc import Callable, Collection, Mapping -from dataclasses import dataclass import json import sys +from collections.abc import Callable, Collection, Mapping +from dataclasses import dataclass from typing import Protocol from alembic.config import Config @@ -23,7 +23,6 @@ load_v2_migration_database_settings, ) - POLICY_MIGRATION_REVISION = "711ec2f0a5a5" V2_ALEMBIC_CONFIG = "alembic-v2.ini" diff --git a/policyengine_api/data/v2/settings.py b/policyengine_api/data/v2/settings.py index 06920ac9c..c169f847c 100644 --- a/policyengine_api/data/v2/settings.py +++ b/policyengine_api/data/v2/settings.py @@ -7,17 +7,15 @@ from __future__ import annotations -from collections.abc import Mapping -from dataclasses import dataclass, field -from functools import lru_cache import os import re -from typing import Callable +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from functools import cache from sqlalchemy.engine import URL, make_url from sqlalchemy.exc import ArgumentError - V2_RUNTIME_DATABASE_URL = "V2_RUNTIME_DATABASE_URL" V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE = "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE" V2_MIGRATION_DATABASE_URL = "V2_MIGRATION_DATABASE_URL" @@ -88,7 +86,7 @@ def _required(environ: Mapping[str, str], name: str) -> str: return value.strip() -@lru_cache(maxsize=None) +@cache def _load_secret_from_secret_manager(resource_name: str) -> str: """Resolve a v2 runtime URL only when a preview request selects it.""" diff --git a/policyengine_api/dataset_display.py b/policyengine_api/dataset_display.py index 8859db018..b303ed3bd 100644 --- a/policyengine_api/dataset_display.py +++ b/policyengine_api/dataset_display.py @@ -2,7 +2,6 @@ from __future__ import annotations - DATASET_FAMILY_DISPLAY_LABELS = ( ("populace_", "Microcosm"), ("enhanced_frs_", "Enhanced FRS"), diff --git a/policyengine_api/extensions.py b/policyengine_api/extensions.py index 064b55d22..d9318fbb1 100644 --- a/policyengine_api/extensions.py +++ b/policyengine_api/extensions.py @@ -2,5 +2,4 @@ from flask_caching import Cache - cache = Cache() diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 41a288836..823f88452 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -13,6 +13,12 @@ from policyengine_api.json_types import JSONObject if TYPE_CHECKING: + from policyengine_api.services.v2.households.types import ( + HouseholdPage, + HouseholdRead, + NativeHouseholdCreation, + NativeHouseholdCreationInput, + ) from policyengine_api.services.v2.metadata.types import ( MetadataCanonicalParameterValue, MetadataDataset, @@ -33,11 +39,11 @@ PolicyPage, PolicyRead, ) - from policyengine_api.services.v2.households.types import ( - HouseholdPage, - HouseholdRead, - NativeHouseholdCreation, - NativeHouseholdCreationInput, + from policyengine_api.services.v2.user_households.types import ( + UserHouseholdCreationInput, + UserHouseholdPage, + UserHouseholdRead, + UserHouseholdUpdateInput, ) from policyengine_api.services.v2.user_policies.types import ( UserPolicyCreationInput, @@ -45,12 +51,6 @@ UserPolicyRead, UserPolicyUpdateInput, ) - from policyengine_api.services.v2.user_households.types import ( - UserHouseholdCreationInput, - UserHouseholdPage, - UserHouseholdRead, - UserHouseholdUpdateInput, - ) class MetadataReader(Protocol): @@ -71,20 +71,20 @@ def list_models( *, offset: int = 0, limit: int = 100, - ) -> "MetadataPageResult[MetadataModel]": ... + ) -> MetadataPageResult[MetadataModel]: ... def get_model( self, country_id: str, model_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataModel]": ... + ) -> MetadataDetailResult[MetadataModel]: ... def get_model_by_country( self, country_id: str, policyengine_version: str | None = None, - ) -> "MetadataModelSelectionResult": ... + ) -> MetadataModelSelectionResult: ... def list_model_versions( self, @@ -93,14 +93,14 @@ def list_model_versions( *, offset: int = 0, limit: int = 100, - ) -> "MetadataPageResult[MetadataModelVersionDetail]": ... + ) -> MetadataPageResult[MetadataModelVersionDetail]: ... def get_model_version( self, country_id: str, version_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataModelVersionDetail]": ... + ) -> MetadataDetailResult[MetadataModelVersionDetail]: ... def list_variables( self, @@ -110,14 +110,14 @@ def list_variables( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> "MetadataPageResult[MetadataVariable]": ... + ) -> MetadataPageResult[MetadataVariable]: ... def get_variable( self, country_id: str, variable_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataVariable]": ... + ) -> MetadataDetailResult[MetadataVariable]: ... def list_parameters( self, @@ -127,14 +127,14 @@ def list_parameters( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> "MetadataPageResult[MetadataParameterSummary]": ... + ) -> MetadataPageResult[MetadataParameterSummary]: ... def get_parameter( self, country_id: str, parameter_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataParameterSummary]": ... + ) -> MetadataDetailResult[MetadataParameterSummary]: ... def list_parameter_children( self, @@ -144,7 +144,7 @@ def list_parameter_children( parent_path: str = "", offset: int = 0, limit: int = 100, - ) -> "MetadataPageResult[MetadataParameterChild]": ... + ) -> MetadataPageResult[MetadataParameterChild]: ... def list_parameter_values( self, @@ -156,14 +156,14 @@ def list_parameter_values( offset: int = 0, limit: int = 100, now: datetime | None = None, - ) -> "MetadataPageResult[MetadataCanonicalParameterValue]": ... + ) -> MetadataPageResult[MetadataCanonicalParameterValue]: ... def get_parameter_value( self, country_id: str, value_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataCanonicalParameterValue]": ... + ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: ... def list_datasets( self, @@ -172,14 +172,14 @@ def list_datasets( *, offset: int = 0, limit: int = 100, - ) -> "MetadataPageResult[MetadataDataset]": ... + ) -> MetadataPageResult[MetadataDataset]: ... def get_dataset( self, country_id: str, dataset_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataDataset]": ... + ) -> MetadataDetailResult[MetadataDataset]: ... def list_regions( self, @@ -189,27 +189,27 @@ def list_regions( region_type: str | None = None, offset: int = 0, limit: int = 100, - ) -> "MetadataPageResult[MetadataRegion]": ... + ) -> MetadataPageResult[MetadataRegion]: ... def get_region( self, country_id: str, region_id: UUID, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataRegion]": ... + ) -> MetadataDetailResult[MetadataRegion]: ... def get_region_by_code( self, country_id: str, region_code: str, policyengine_version: str | None = None, - ) -> "MetadataDetailResult[MetadataRegion]": ... + ) -> MetadataDetailResult[MetadataRegion]: ... def get_economy_options( self, country_id: str, policyengine_version: str | None = None, - ) -> "MetadataEconomyOptionsResult": ... + ) -> MetadataEconomyOptionsResult: ... class V2PolicyResourceService(Protocol): @@ -217,10 +217,10 @@ class V2PolicyResourceService(Protocol): def create_policy( self, - command: "NativePolicyCreationInput", - ) -> "NativePolicyCreation": ... + command: NativePolicyCreationInput, + ) -> NativePolicyCreation: ... - def get_policy(self, *, country_id: str, policy_id: UUID) -> "PolicyRead": ... + def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: ... def list_policies( self, @@ -229,7 +229,7 @@ def list_policies( tax_benefit_model_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> "PolicyPage": ... + ) -> PolicyPage: ... class V2HouseholdResourceService(Protocol): @@ -237,15 +237,15 @@ class V2HouseholdResourceService(Protocol): def create_household( self, - household_input: "NativeHouseholdCreationInput", - ) -> "NativeHouseholdCreation": ... + household_input: NativeHouseholdCreationInput, + ) -> NativeHouseholdCreation: ... def get_household( self, *, country_id: str, household_id: UUID, - ) -> "HouseholdRead": ... + ) -> HouseholdRead: ... def list_households( self, @@ -254,7 +254,7 @@ def list_households( default_year: int | None, offset: int, limit: int, - ) -> "HouseholdPage": ... + ) -> HouseholdPage: ... class V2UserPolicyResourceService(Protocol): @@ -262,15 +262,15 @@ class V2UserPolicyResourceService(Protocol): def create_user_policy( self, - association_input: "UserPolicyCreationInput", - ) -> "UserPolicyRead": ... + association_input: UserPolicyCreationInput, + ) -> UserPolicyRead: ... def get_user_policy( self, *, country_id: str, association_id: UUID, - ) -> "UserPolicyRead": ... + ) -> UserPolicyRead: ... def list_user_policies( self, @@ -280,15 +280,15 @@ def list_user_policies( policy_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> "UserPolicyPage": ... + ) -> UserPolicyPage: ... def patch_user_policy( self, *, country_id: str, association_id: UUID, - association_input: "UserPolicyUpdateInput", - ) -> "UserPolicyRead": ... + association_input: UserPolicyUpdateInput, + ) -> UserPolicyRead: ... def delete_user_policy( self, @@ -303,15 +303,15 @@ class V2UserHouseholdResourceService(Protocol): def create_user_household( self, - association_input: "UserHouseholdCreationInput", - ) -> "UserHouseholdRead": ... + association_input: UserHouseholdCreationInput, + ) -> UserHouseholdRead: ... def get_user_household( self, *, country_id: str, association_id: UUID, - ) -> "UserHouseholdRead": ... + ) -> UserHouseholdRead: ... def list_user_households( self, @@ -321,15 +321,15 @@ def list_user_households( household_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> "UserHouseholdPage": ... + ) -> UserHouseholdPage: ... def patch_user_household( self, *, country_id: str, association_id: UUID, - association_input: "UserHouseholdUpdateInput", - ) -> "UserHouseholdRead": ... + association_input: UserHouseholdUpdateInput, + ) -> UserHouseholdRead: ... def delete_user_household( self, @@ -457,7 +457,7 @@ class NativeRouteDependencies: ) = None @classmethod - def defaults(cls) -> "NativeRouteDependencies": + def defaults(cls) -> NativeRouteDependencies: """Build lazy defaults without importing country models unnecessarily.""" return cls( readiness_probe=_default_readiness_probe, diff --git a/policyengine_api/fastapi_routes/health.py b/policyengine_api/fastapi_routes/health.py index 9ea431702..8b7eb33e3 100644 --- a/policyengine_api/fastapi_routes/health.py +++ b/policyengine_api/fastapi_routes/health.py @@ -5,10 +5,11 @@ from typing import Literal from fastapi import APIRouter, HTTPException -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies from pydantic import BaseModel from starlette.responses import PlainTextResponse +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies + class HealthResponse(BaseModel): status: Literal["healthy"] diff --git a/policyengine_api/fastapi_routes/metadata.py b/policyengine_api/fastapi_routes/metadata.py index 40dcbcf10..ee9c928c8 100644 --- a/policyengine_api/fastapi_routes/metadata.py +++ b/policyengine_api/fastapi_routes/metadata.py @@ -5,6 +5,8 @@ from typing import Literal, TypedDict from fastapi import APIRouter +from starlette.responses import Response, StreamingResponse + from policyengine_api.country_validation import ( InvalidCountryError, ensure_supported_country, @@ -19,7 +21,6 @@ iter_body_chunks, should_stream_body, ) -from starlette.responses import Response, StreamingResponse class MetadataSuccessPayload(TypedDict): diff --git a/policyengine_api/fastapi_routes/query_parameters.py b/policyengine_api/fastapi_routes/query_parameters.py index 98670d410..608406bca 100644 --- a/policyengine_api/fastapi_routes/query_parameters.py +++ b/policyengine_api/fastapi_routes/query_parameters.py @@ -15,7 +15,6 @@ validate_scalar_query_multiplicity, ) - QueryParametersT = TypeVar("QueryParametersT", bound=StrictQueryParameters) diff --git a/policyengine_api/fastapi_routes/responses.py b/policyengine_api/fastapi_routes/responses.py index 07c3e57a5..7c1d160de 100644 --- a/policyengine_api/fastapi_routes/responses.py +++ b/policyengine_api/fastapi_routes/responses.py @@ -4,9 +4,10 @@ import json -from policyengine_api.json_types import JSONValue from starlette.responses import Response +from policyengine_api.json_types import JSONValue + def render_legacy_json(content: JSONValue) -> bytes: """Serialize once with the same JSON encoder used by legacy metadata.""" diff --git a/policyengine_api/fastapi_routes/specification.py b/policyengine_api/fastapi_routes/specification.py index 8d45bf32e..f104159f6 100644 --- a/policyengine_api/fastapi_routes/specification.py +++ b/policyengine_api/fastapi_routes/specification.py @@ -3,9 +3,10 @@ from __future__ import annotations from fastapi import APIRouter -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies from starlette.responses import JSONResponse +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies + def build_specification_router( dependencies: NativeRouteDependencies, diff --git a/policyengine_api/fastapi_routes/v2/households/document_models.py b/policyengine_api/fastapi_routes/v2/households/document_models.py index a7ab499a8..25a2d4dbc 100644 --- a/policyengine_api/fastapi_routes/v2/households/document_models.py +++ b/policyengine_api/fastapi_routes/v2/households/document_models.py @@ -10,14 +10,13 @@ Field, JsonValue, StringConstraints, - model_serializer, field_validator, + model_serializer, ) from pydantic.json_schema import SkipJsonSchema from policyengine_api.spm import SPMSelection - DocumentIdentifier = Annotated[ str, StringConstraints(strip_whitespace=False, min_length=1, max_length=255), diff --git a/policyengine_api/fastapi_routes/v2/households/request_models.py b/policyengine_api/fastapi_routes/v2/households/request_models.py index 31ab80e4f..32619006a 100644 --- a/policyengine_api/fastapi_routes/v2/households/request_models.py +++ b/policyengine_api/fastapi_routes/v2/households/request_models.py @@ -12,7 +12,6 @@ from policyengine_api.query_parameters import CountryId, DefaultYear from policyengine_api.services.v2.households.types import NativeHouseholdCreationInput - MAXIMUM_HOUSEHOLD_REQUEST_BYTES = 1_048_576 @@ -26,7 +25,7 @@ class HouseholdCreateRequest(BaseModel): household_data: HouseholdDocument @model_validator(mode="after") - def require_country_document(self) -> "HouseholdCreateRequest": + def require_country_document(self) -> HouseholdCreateRequest: expected_type = ( USHouseholdDocument if self.country_id == "us" else UKHouseholdDocument ) diff --git a/policyengine_api/fastapi_routes/v2/households/response_models.py b/policyengine_api/fastapi_routes/v2/households/response_models.py index fffb02f4a..4388a4765 100644 --- a/policyengine_api/fastapi_routes/v2/households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/households/response_models.py @@ -33,7 +33,7 @@ class HouseholdItem(StrictHouseholdAPIModel): updated_at: datetime @classmethod - def from_read(cls, item: HouseholdRead) -> "HouseholdItem": + def from_read(cls, item: HouseholdRead) -> HouseholdItem: return cls.model_validate(item) @@ -48,7 +48,7 @@ class HouseholdPageResult(StrictHouseholdAPIModel): has_more: bool @classmethod - def from_page(cls, page: HouseholdPage) -> "HouseholdPageResult": + def from_page(cls, page: HouseholdPage) -> HouseholdPageResult: return cls( items=[HouseholdItem.from_read(item) for item in page.items], offset=page.offset, diff --git a/policyengine_api/fastapi_routes/v2/metadata/common.py b/policyengine_api/fastapi_routes/v2/metadata/common.py index fefc8b0a9..cb483c011 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/common.py +++ b/policyengine_api/fastapi_routes/v2/metadata/common.py @@ -14,21 +14,20 @@ MetadataCatalogVersionNotFoundError, UnsupportedPreviewCountryError, ) -from policyengine_api.services.v2.metadata.validators import ( - InvalidMetadataPageError, - MetadataResourceNotFoundError, +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2MetadataResourceReader, ) from policyengine_api.fastapi_routes.v2.errors import ( V2ErrorResponse, v2_error_response, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.fastapi_routes.dependencies import ( - NativeRouteDependencies, - V2MetadataResourceReader, +from policyengine_api.services.v2.metadata.validators import ( + InvalidMetadataPageError, + MetadataResourceNotFoundError, ) - ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { 400: { "model": V2ErrorResponse, diff --git a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py index e3da1f81e..ed8a0aca2 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py @@ -8,7 +8,11 @@ from fastapi import APIRouter, Query from starlette.responses import JSONResponse -from policyengine_api.services.v2.metadata.types import MetadataRegionType +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.fastapi_routes.v2.metadata.common import ( + ERROR_RESPONSES, + read_resource, +) from policyengine_api.fastapi_routes.v2.metadata.response_models import ( MetadataDatasetDetailResponse, MetadataDatasetPageResponse, @@ -16,12 +20,7 @@ MetadataRegionDetailResponse, MetadataRegionPageResponse, ) -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.fastapi_routes.v2.metadata.common import ( - ERROR_RESPONSES, - read_resource, -) - +from policyengine_api.services.v2.metadata.types import MetadataRegionType Offset = Annotated[int, Query(ge=0)] Limit = Annotated[int, Query(ge=1, le=500)] diff --git a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py index 01693f512..5d7315641 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py @@ -8,6 +8,11 @@ from fastapi import APIRouter, Query from starlette.responses import JSONResponse +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.fastapi_routes.v2.metadata.common import ( + ERROR_RESPONSES, + read_resource, +) from policyengine_api.fastapi_routes.v2.metadata.response_models import ( MetadataModelDetailResponse, MetadataModelPageResponse, @@ -17,12 +22,6 @@ MetadataVariableDetailResponse, MetadataVariablePageResponse, ) -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.fastapi_routes.v2.metadata.common import ( - ERROR_RESPONSES, - read_resource, -) - Offset = Annotated[int, Query(ge=0)] Limit = Annotated[int, Query(ge=1, le=500)] diff --git a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py index 714c3e28d..5c7f19f4f 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py @@ -8,6 +8,11 @@ from fastapi import APIRouter, Query from starlette.responses import JSONResponse +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.fastapi_routes.v2.metadata.common import ( + ERROR_RESPONSES, + read_resource, +) from policyengine_api.fastapi_routes.v2.metadata.response_models import ( MetadataParameterChildPageResponse, MetadataParameterDetailResponse, @@ -15,12 +20,6 @@ MetadataParameterValueDetailResponse, MetadataParameterValuePageResponse, ) -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.fastapi_routes.v2.metadata.common import ( - ERROR_RESPONSES, - read_resource, -) - Offset = Annotated[int, Query(ge=0)] Limit = Annotated[int, Query(ge=1, le=500)] diff --git a/policyengine_api/fastapi_routes/v2/metadata/response_models.py b/policyengine_api/fastapi_routes/v2/metadata/response_models.py index 8d6d965a9..6637998e9 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/response_models.py +++ b/policyengine_api/fastapi_routes/v2/metadata/response_models.py @@ -20,7 +20,6 @@ StrictResponseModel, ) - ResourceT = TypeVar("ResourceT") diff --git a/policyengine_api/fastapi_routes/v2/policies/request_models.py b/policyengine_api/fastapi_routes/v2/policies/request_models.py index e3180b94f..75aa5ebc0 100644 --- a/policyengine_api/fastapi_routes/v2/policies/request_models.py +++ b/policyengine_api/fastapi_routes/v2/policies/request_models.py @@ -2,7 +2,6 @@ from policyengine_api.services.v2.policies.types import PolicyCreationInput - MAXIMUM_POLICY_REQUEST_BYTES = 1_048_576 diff --git a/policyengine_api/fastapi_routes/v2/policies/response_models.py b/policyengine_api/fastapi_routes/v2/policies/response_models.py index 94784c831..ec44894fd 100644 --- a/policyengine_api/fastapi_routes/v2/policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/policies/response_models.py @@ -40,7 +40,7 @@ class PolicyItem(StrictPolicyAPIModel): parameter_values: list[PolicyParameterValueItem] @classmethod - def from_read(cls, item: PolicyRead) -> "PolicyItem": + def from_read(cls, item: PolicyRead) -> PolicyItem: return cls.model_validate(item) @@ -55,7 +55,7 @@ class PolicyPageResult(StrictPolicyAPIModel): has_more: bool @classmethod - def from_page(cls, page: PolicyPage) -> "PolicyPageResult": + def from_page(cls, page: PolicyPage) -> PolicyPageResult: return cls( items=[PolicyItem.from_read(item) for item in page.items], offset=page.offset, diff --git a/policyengine_api/fastapi_routes/v2/policies/routes.py b/policyengine_api/fastapi_routes/v2/policies/routes.py index 8e2a5d754..383c4c0f2 100644 --- a/policyengine_api/fastapi_routes/v2/policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/policies/routes.py @@ -14,6 +14,16 @@ MetadataCatalogUnavailableError, MetadataCatalogVersionNotFoundError, ) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2PolicyResourceService, +) +from policyengine_api.fastapi_routes.query_parameters import query_dependency +from policyengine_api.fastapi_routes.v2.errors import ( + V2RequestTooLargeError, + v2_error_response, +) from policyengine_api.fastapi_routes.v2.policies.request_models import ( MAXIMUM_POLICY_REQUEST_BYTES, PolicyCreateRequest, @@ -26,9 +36,10 @@ PolicyPageResponse, PolicyPageResult, ) -from policyengine_api.fastapi_routes.v2.errors import ( - V2RequestTooLargeError, - v2_error_response, +from policyengine_api.query_parameters import ( + PolicyCollectionQuery, + PolicyCreateQuery, + PolicyDetailQuery, ) from policyengine_api.services.v2.policies.types import NativePolicyCreationInput from policyengine_api.services.v2.policies.validators import ( @@ -37,17 +48,6 @@ PolicyCreationIntegrityError, PolicyNotFoundError, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.fastapi_routes.dependencies import ( - NativeRouteDependencies, - V2PolicyResourceService, -) -from policyengine_api.fastapi_routes.query_parameters import query_dependency -from policyengine_api.query_parameters import ( - PolicyCollectionQuery, - PolicyCreateQuery, - PolicyDetailQuery, -) async def enforce_policy_request_size(request: Request) -> None: diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index fc2624a2d..d5b97d3c8 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -7,6 +7,9 @@ from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse +from policyengine_api.fastapi_routes.v2.households.routes import ( + build_v2_household_router, +) from policyengine_api.fastapi_routes.v2.metadata.geography_routes import ( build_v2_metadata_geography_router, ) @@ -16,16 +19,13 @@ from policyengine_api.fastapi_routes.v2.metadata.parameter_routes import ( build_v2_metadata_parameter_router, ) -from policyengine_api.fastapi_routes.v2.households.routes import ( - build_v2_household_router, -) from policyengine_api.fastapi_routes.v2.policies.routes import build_v2_policy_router -from policyengine_api.fastapi_routes.v2.user_policies.routes import ( - build_v2_user_policy_router, -) from policyengine_api.fastapi_routes.v2.user_households.routes import ( build_v2_user_household_router, ) +from policyengine_api.fastapi_routes.v2.user_policies.routes import ( + build_v2_user_policy_router, +) def build_v2_router( diff --git a/policyengine_api/fastapi_routes/v2/user_households/response_models.py b/policyengine_api/fastapi_routes/v2/user_households/response_models.py index d6ba2fd01..ccb53cd0a 100644 --- a/policyengine_api/fastapi_routes/v2/user_households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_households/response_models.py @@ -31,7 +31,7 @@ class UserHouseholdItem(StrictUserHouseholdAPIModel): updated_at: datetime @classmethod - def from_read(cls, item: UserHouseholdRead) -> "UserHouseholdItem": + def from_read(cls, item: UserHouseholdRead) -> UserHouseholdItem: return cls.model_validate(item) @@ -46,7 +46,7 @@ class UserHouseholdPageResult(StrictUserHouseholdAPIModel): has_more: bool @classmethod - def from_page(cls, page: UserHouseholdPage) -> "UserHouseholdPageResult": + def from_page(cls, page: UserHouseholdPage) -> UserHouseholdPageResult: return cls( items=[UserHouseholdItem.from_read(item) for item in page.items], offset=page.offset, diff --git a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py index caef44df3..da0c9ed21 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py @@ -9,11 +9,11 @@ from pydantic import BaseModel, ConfigDict from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse +from policyengine_api.query_parameters import CountryId, ResourceId, UserId from policyengine_api.services.v2.user_policies.types import ( UserPolicyPage, UserPolicyRead, ) -from policyengine_api.query_parameters import CountryId, ResourceId, UserId class StrictUserPolicyAPIModel(BaseModel): @@ -33,7 +33,7 @@ class UserPolicyItem(StrictUserPolicyAPIModel): updated_at: datetime @classmethod - def from_read(cls, item: UserPolicyRead) -> "UserPolicyItem": + def from_read(cls, item: UserPolicyRead) -> UserPolicyItem: return cls.model_validate(item) @@ -48,7 +48,7 @@ class UserPolicyPageResult(StrictUserPolicyAPIModel): has_more: bool @classmethod - def from_page(cls, page: UserPolicyPage) -> "UserPolicyPageResult": + def from_page(cls, page: UserPolicyPage) -> UserPolicyPageResult: return cls( items=[UserPolicyItem.from_read(item) for item in page.items], offset=page.offset, diff --git a/policyengine_api/fastapi_routes/v2/user_policies/routes.py b/policyengine_api/fastapi_routes/v2/user_policies/routes.py index 7b4f78c77..2be0184d0 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/routes.py @@ -11,6 +11,12 @@ from starlette.responses import JSONResponse, Response from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2UserPolicyResourceService, +) +from policyengine_api.fastapi_routes.query_parameters import query_dependency +from policyengine_api.fastapi_routes.v2.errors import v2_error_response from policyengine_api.fastapi_routes.v2.user_policies.request_models import ( UserPolicyCreateRequest, UserPolicyPatchRequest, @@ -23,21 +29,15 @@ UserPolicyPageResponse, UserPolicyPageResult, ) -from policyengine_api.fastapi_routes.v2.errors import v2_error_response +from policyengine_api.query_parameters import ( + CountryQuery, + UserPolicyCollectionQuery, +) from policyengine_api.services.v2.user_policies.validators import ( - UserPolicyNotFoundError, AssociationCountryConflictError, AssociationPolicyNotFoundError, AssociationUserNotFoundError, -) -from policyengine_api.fastapi_routes.dependencies import ( - NativeRouteDependencies, - V2UserPolicyResourceService, -) -from policyengine_api.fastapi_routes.query_parameters import query_dependency -from policyengine_api.query_parameters import ( - CountryQuery, - UserPolicyCollectionQuery, + UserPolicyNotFoundError, ) diff --git a/policyengine_api/json_types.py b/policyengine_api/json_types.py index dc6f4e97f..c30e6037c 100644 --- a/policyengine_api/json_types.py +++ b/policyengine_api/json_types.py @@ -4,7 +4,6 @@ from typing import TypeAlias - JSONScalar: TypeAlias = str | int | float | bool | None JSONValue: TypeAlias = JSONScalar | list["JSONValue"] | dict[str, "JSONValue"] JSONObject: TypeAlias = dict[str, JSONValue] diff --git a/policyengine_api/libs/gateway_auth.py b/policyengine_api/libs/gateway_auth.py index 34e2120fe..808202e5b 100644 --- a/policyengine_api/libs/gateway_auth.py +++ b/policyengine_api/libs/gateway_auth.py @@ -12,12 +12,10 @@ import os import threading import time -from functools import lru_cache -from typing import Optional +from functools import cache import httpx - GATEWAY_AUTH_ISSUER_ENV = "GATEWAY_AUTH_ISSUER" GATEWAY_AUTH_AUDIENCE_ENV = "GATEWAY_AUTH_AUDIENCE" GATEWAY_AUTH_CLIENT_ID_ENV = "GATEWAY_AUTH_CLIENT_ID" @@ -56,7 +54,7 @@ def gateway_auth_required() -> bool: ) -@lru_cache(maxsize=None) +@cache def _load_secret_from_secret_manager(resource_name: str) -> str: """Fetch one secret payload from Google Secret Manager.""" from google.cloud import secretmanager @@ -110,11 +108,11 @@ class GatewayAuthTokenProvider: def __init__( self, - issuer: Optional[str] = None, - audience: Optional[str] = None, - client_id: Optional[str] = None, - client_secret: Optional[str] = None, - client_secret_resource: Optional[str] = None, + issuer: str | None = None, + audience: str | None = None, + client_id: str | None = None, + client_secret: str | None = None, + client_secret_resource: str | None = None, *, http_timeout: float = 10.0, ): @@ -144,7 +142,7 @@ def __init__( else os.environ.get(GATEWAY_AUTH_CLIENT_SECRET_RESOURCE_ENV, "") ) self._http_timeout = http_timeout - self._token: Optional[str] = None + self._token: str | None = None self._expires_at: float = 0.0 self._lock = threading.Lock() diff --git a/policyengine_api/libs/simulation_api.py b/policyengine_api/libs/simulation_api.py index 8d95ceee1..a3aa22d15 100644 --- a/policyengine_api/libs/simulation_api.py +++ b/policyengine_api/libs/simulation_api.py @@ -7,8 +7,8 @@ ModalBudgetWindowBatchExecution, ModalSimulationExecution, SimulationAPIClient, - SimulationEntrypointClient, SimulationAPIModal, + SimulationEntrypointClient, resolve_simulation_api_url, resolve_simulation_entrypoint_url, simulation_api, diff --git a/policyengine_api/libs/simulation_api_modal.py b/policyengine_api/libs/simulation_api_modal.py index 3632c610b..2c472af46 100644 --- a/policyengine_api/libs/simulation_api_modal.py +++ b/policyengine_api/libs/simulation_api_modal.py @@ -4,8 +4,8 @@ ModalBudgetWindowBatchExecution, ModalSimulationExecution, SimulationAPIClient, - SimulationEntrypointClient, SimulationAPIModal, + SimulationEntrypointClient, resolve_simulation_api_url, resolve_simulation_entrypoint_url, simulation_api, diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index aa2a2bd17..997cb0352 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -3,11 +3,11 @@ import os import sys from dataclasses import dataclass, field -from typing import Optional from urllib.parse import urlparse import httpx from policyengine_observability import instrument_httpx + from policyengine_api.gcp_logging import logger from policyengine_api.libs.gateway_auth import ( GatewayAuthError, @@ -18,13 +18,13 @@ ) from policyengine_api.migration_flags import get_sim_entrypoint from policyengine_api.observability import get_runtime +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.request_context import ( REQUEST_ID_HEADER, current_observability_id, current_request_id, ) -from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER -from policyengine_api.worker_spm import validate_worker_spm, raise_worker_spm_error +from policyengine_api.worker_spm import raise_worker_spm_error, validate_worker_spm def _required_base_url(env_name: str) -> str: @@ -77,10 +77,10 @@ class ModalSimulationExecution: job_id: str status: str - result: Optional[dict] = None - error: Optional[str] = None - policyengine_bundle: Optional[dict] = None - resolved_app_name: Optional[str] = None + result: dict | None = None + error: str | None = None + policyengine_bundle: dict | None = None + resolved_app_name: str | None = None @property def name(self) -> str: @@ -96,13 +96,13 @@ class ModalBudgetWindowBatchExecution: batch_job_id: str status: str - progress: Optional[int] = None + progress: int | None = None completed_years: list[str] = field(default_factory=list) running_years: list[str] = field(default_factory=list) queued_years: list[str] = field(default_factory=list) failed_years: list[str] = field(default_factory=list) - result: Optional[dict] = None - error: Optional[str] = None + result: dict | None = None + error: str | None = None @property def name(self) -> str: @@ -155,8 +155,8 @@ def _normalize_submission_payload(self, payload: dict) -> dict: raise ValueError("Dataset overrides are not supported") from policyengine_api.constants import ( - POLICYENGINE_VERSION, COUNTRY_PACKAGE_VERSIONS, + POLICYENGINE_VERSION, ) country = payload.get("country", "us") @@ -367,8 +367,8 @@ def get_spm_capability(self, country, version=None, *, policyengine_version=None def resolve_app_name( self, country: str, - version: Optional[str] = None, - policyengine_version: Optional[str] = None, + version: str | None = None, + policyengine_version: str | None = None, ) -> tuple[str, str]: """Resolve the current gateway app name for a country/model version.""" if policyengine_version is not None: @@ -524,9 +524,7 @@ def get_execution_status(self, execution: ModalSimulationExecution) -> str: """ return execution.status - def get_execution_result( - self, execution: ModalSimulationExecution - ) -> Optional[dict]: + def get_execution_result(self, execution: ModalSimulationExecution) -> dict | None: """ Get the result from a completed execution. diff --git a/policyengine_api/migration_flags.py b/policyengine_api/migration_flags.py index d4496e6d5..34843a1c2 100644 --- a/policyengine_api/migration_flags.py +++ b/policyengine_api/migration_flags.py @@ -128,7 +128,7 @@ class RouteImplementationSettings: metadata: RouteImplementation @classmethod - def from_environment(cls) -> "RouteImplementationSettings": + def from_environment(cls) -> RouteImplementationSettings: """Read and validate each Stage 6 route selector independently.""" return cls( health=get_route_impl("health"), diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index e4432ace1..cef7e8985 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -6,21 +6,21 @@ import flask from policyengine_observability import ObservabilityRuntime + from policyengine_api.gcp_logging import logger from policyengine_api.migration_flags import ( RouteImplementation, get_migration_log_context, infer_route_group, ) -from policyengine_api.request_context import ( - REQUEST_ID_HEADER, - generate_request_id, -) from policyengine_api.observability.identifiers import ( OBSERVABILITY_ID_HEADER, normalize_observability_id, ) - +from policyengine_api.request_context import ( + REQUEST_ID_HEADER, + generate_request_id, +) V2_METADATA_RESOURCE_SEGMENTS = frozenset( { diff --git a/policyengine_api/observability/runtime.py b/policyengine_api/observability/runtime.py index eefd45185..aea6f10c2 100644 --- a/policyengine_api/observability/runtime.py +++ b/policyengine_api/observability/runtime.py @@ -18,7 +18,6 @@ process_instance_id, ) - DISPATCH_ATTRIBUTE_KEYS = frozenset({"observability_id"}) diff --git a/policyengine_api/observability/stages.py b/policyengine_api/observability/stages.py index eaba1b7d9..efa020b1d 100644 --- a/policyengine_api/observability/stages.py +++ b/policyengine_api/observability/stages.py @@ -2,10 +2,10 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from enum import StrEnum from types import MappingProxyType -from typing import Mapping class RunConfiguration(StrEnum): diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 3be792058..3cbb45202 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping import json +from collections.abc import Iterable, Mapping from types import UnionType from typing import Annotated, Any, Literal, TypeVar, Union, get_args, get_origin from uuid import UUID @@ -26,7 +26,6 @@ BUDGET_WINDOW_MAX_YEARS, ) - DEFAULT_QUERY_LIMIT = 100 MAXIMUM_QUERY_LIMIT = 500 MAXIMUM_USER_ID_LENGTH = 255 diff --git a/policyengine_api/request_context.py b/policyengine_api/request_context.py index 08c88d1d2..6e7fd8ef3 100644 --- a/policyengine_api/request_context.py +++ b/policyengine_api/request_context.py @@ -6,6 +6,7 @@ from contextvars import ContextVar import flask + from policyengine_api.observability.identifiers import ( generate_observability_id, normalize_observability_id, diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 858650511..60837fa7d 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -1,27 +1,29 @@ +import json +from http import HTTPStatus + from flask import Blueprint, Response, request from pydantic import ValidationError + +from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.query_parameters import ( - AnnualEconomyQuery, BUDGET_WINDOW_MAX_YEARS, + AnnualEconomyQuery, BudgetWindowEconomyQuery, parse_multidict_query, ) +from policyengine_api.response_factory import _make_error_response from policyengine_api.services.economy_service import ( - EconomyService, - EconomicImpactResult, BudgetWindowEconomicImpactResult, + EconomicImpactResult, + EconomyService, ) -from policyengine_api.response_factory import _make_error_response -from policyengine_api.utils import get_current_law_policy_id -from policyengine_api.utils.payload_validators import validate_country -from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.spm import ( SPMValidationError, readable_validation_error, spm_error_detail, ) -import json -from http import HTTPStatus +from policyengine_api.utils import get_current_law_policy_id +from policyengine_api.utils.payload_validators import validate_country economy_bp = Blueprint("economy", __name__) economy_service = EconomyService() diff --git a/policyengine_api/routes/home_routes.py b/policyengine_api/routes/home_routes.py index fdc230705..ac05ccf44 100644 --- a/policyengine_api/routes/home_routes.py +++ b/policyengine_api/routes/home_routes.py @@ -1,6 +1,5 @@ from flask import Blueprint - home_bp = Blueprint("home", __name__) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 300f2d9d3..13c8888b4 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -21,16 +21,16 @@ start_observability_id, ) from policyengine_api.response_factory import _make_error_response -from policyengine_api.services.household_mirroring import ( - HouseholdMirrorUnavailableError, - process_household_event_after_commit, -) from policyengine_api.services.household_calculation_service import ( HouseholdCalculationService, HouseholdNotFoundError, InvalidHouseholdInputsError, PolicyNotFoundError, ) +from policyengine_api.services.household_mirroring import ( + HouseholdMirrorUnavailableError, + process_household_event_after_commit, +) from policyengine_api.services.household_service import ( HouseholdPersistenceError, HouseholdService, @@ -258,18 +258,17 @@ def get_household(country_id: str, household_id: int) -> Response: result = None if household is None else _serialize_household(household) if result is None: raise NotFound(f"Household #{household_id} not found.") - else: - return Response( - json.dumps( - { - "status": "ok", - "message": None, - "result": result, - } - ), - status=200, - mimetype="application/json", - ) + return Response( + json.dumps( + { + "status": "ok", + "message": None, + "result": result, + } + ), + status=200, + mimetype="application/json", + ) @household_bp.route("//household", methods=["POST"]) diff --git a/policyengine_api/routes/metadata_routes.py b/policyengine_api/routes/metadata_routes.py index cc1c310d0..274585fe4 100644 --- a/policyengine_api/routes/metadata_routes.py +++ b/policyengine_api/routes/metadata_routes.py @@ -1,12 +1,13 @@ import json + from flask import Blueprint, Response +from policyengine_api.services.metadata_service import MetadataService from policyengine_api.utils.payload_validators import validate_country from policyengine_api.utils.streaming_json import ( iter_body_chunks, should_stream_body, ) -from policyengine_api.services.metadata_service import MetadataService metadata_bp = Blueprint("metadata", __name__) metadata_service = MetadataService() diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 3918703e7..d27884f2a 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -20,15 +20,15 @@ mirror_policy_after_commit, ) from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_policy_mirroring import ( + UserPolicyMirrorUnavailableError, + mirror_pending_user_policy_events_after_commit, +) from policyengine_api.services.user_policy_service import ( USER_POLICY_MUTABLE_FIELDS, UserPolicyPersistenceError, UserPolicyService, ) -from policyengine_api.services.user_policy_mirroring import ( - UserPolicyMirrorUnavailableError, - mirror_pending_user_policy_events_after_commit, -) from policyengine_api.utils.payload_validators import ( validate_country, validate_set_policy_payload, diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py index a44f6be56..c3a9d2921 100644 --- a/policyengine_api/routes/reform_impact_routes.py +++ b/policyengine_api/routes/reform_impact_routes.py @@ -1,5 +1,5 @@ -from datetime import datetime import json +from datetime import datetime from typing import Any from flask import Blueprint, Response, request @@ -7,7 +7,6 @@ from policyengine_api.runtime_cache.reform_impacts import CachedReformImpact from policyengine_api.services.reform_impacts_service import ReformImpactsService - reform_impact_bp = Blueprint("reform_impact", __name__) reform_impacts_service = ReformImpactsService() diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index 2b2bc21e7..b105b33ba 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -1,16 +1,16 @@ -from flask import Blueprint, Response, current_app, request -from werkzeug.exceptions import HTTPException, NotFound, BadRequest import json import jsonschema import pydantic +from flask import Blueprint, Response, current_app, request +from werkzeug.exceptions import BadRequest, HTTPException, NotFound +from policyengine_api.constants import CURRENT_YEAR +from policyengine_api.data.v1_models import ReportOutput from policyengine_api.services.report_output_service import ( ReportOutputService, ReportOutputView, ) -from policyengine_api.constants import CURRENT_YEAR -from policyengine_api.data.v1_models import ReportOutput from policyengine_api.utils.payload_validators import validate_country report_output_bp = Blueprint("report_output", __name__) diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index 05275f152..5e9b329e6 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -1,9 +1,9 @@ -from flask import Blueprint, Response, current_app, request -from werkzeug.exceptions import HTTPException, NotFound, BadRequest import json import jsonschema import pydantic +from flask import Blueprint, Response, current_app, request +from werkzeug.exceptions import BadRequest, HTTPException, NotFound from policyengine_api.data.v1_models import Simulation from policyengine_api.response_factory import _make_error_response diff --git a/policyengine_api/routes/system_routes.py b/policyengine_api/routes/system_routes.py index f0e2748c0..450459083 100644 --- a/policyengine_api/routes/system_routes.py +++ b/policyengine_api/routes/system_routes.py @@ -3,7 +3,6 @@ from policyengine_api.readiness import is_ready from policyengine_api.specification import OPENAPI_SPECIFICATION - system_bp = Blueprint("system", __name__) diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index 77175feff..8e36f9707 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,9 +1,11 @@ +import json + from flask import Blueprint, Response, request +from werkzeug.exceptions import BadRequest, NotFound + from policyengine_api.data.v1_models import UserProfile -from policyengine_api.utils.payload_validators import validate_country -import json from policyengine_api.services.user_service import UserService -from werkzeug.exceptions import BadRequest, NotFound +from policyengine_api.utils.payload_validators import validate_country user_profile_bp = Blueprint("user_profile", __name__) user_service = UserService() diff --git a/policyengine_api/runtime_cache/claims.py b/policyengine_api/runtime_cache/claims.py index 6f3c48b70..3fcdc58ab 100644 --- a/policyengine_api/runtime_cache/claims.py +++ b/policyengine_api/runtime_cache/claims.py @@ -10,7 +10,6 @@ record_cache_event, ) - COMPARE_AND_DELETE = """ if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) diff --git a/policyengine_api/runtime_cache/client.py b/policyengine_api/runtime_cache/client.py index 24c1d39f9..e6bbf0dae 100644 --- a/policyengine_api/runtime_cache/client.py +++ b/policyengine_api/runtime_cache/client.py @@ -13,7 +13,6 @@ load_runtime_cache_settings, ) - _clients: dict[int, redis.Redis] = {} _client_lock = Lock() diff --git a/policyengine_api/runtime_cache/core.py b/policyengine_api/runtime_cache/core.py index 0422a158c..ccb1a0df3 100644 --- a/policyengine_api/runtime_cache/core.py +++ b/policyengine_api/runtime_cache/core.py @@ -2,16 +2,16 @@ from __future__ import annotations -from dataclasses import dataclass import hashlib import json import random import time -from typing import Any, Callable, Protocol +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol from policyengine_api.gcp_logging import logger - CACHE_KEY_ROOT = "policyengine" COMPLETED_RESULT_TTL_JITTER_FRACTION = 0.10 diff --git a/policyengine_api/runtime_cache/fake.py b/policyengine_api/runtime_cache/fake.py index 044a07713..11c4e212f 100644 --- a/policyengine_api/runtime_cache/fake.py +++ b/policyengine_api/runtime_cache/fake.py @@ -4,7 +4,7 @@ from collections.abc import Callable from threading import RLock -from typing import Any +from typing import Any, Self class InMemoryCacheBackend: @@ -129,7 +129,7 @@ def eval( return 0 return self.delete(key) - def pipeline(self, transaction: bool = True) -> "InMemoryPipeline": + def pipeline(self, transaction: bool = True) -> InMemoryPipeline: return InMemoryPipeline(self, transaction=transaction) @@ -140,7 +140,7 @@ def __init__(self, backend: InMemoryCacheBackend, *, transaction: bool) -> None: self.operations: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = [] def __getattr__(self, name: str): - def queue(*args: Any, **kwargs: Any) -> "InMemoryPipeline": + def queue(*args: Any, **kwargs: Any) -> InMemoryPipeline: self.operations.append((name, args, kwargs)) return self @@ -156,10 +156,10 @@ def execute(self) -> list[Any]: self.operations.clear() return results - def __enter__(self) -> "InMemoryPipeline": + def __enter__(self) -> Self: return self - def __exit__(self, *_args: Any) -> None: + def __exit__(self, *_args: object) -> None: self.operations.clear() diff --git a/policyengine_api/runtime_cache/household_calculations.py b/policyengine_api/runtime_cache/household_calculations.py index 4ab7ceaf7..f490c3bd2 100644 --- a/policyengine_api/runtime_cache/household_calculations.py +++ b/policyengine_api/runtime_cache/household_calculations.py @@ -15,7 +15,6 @@ validate_spm_calculation_provenance, ) - HOUSEHOLD_CALCULATION_SCHEMA_VERSION = 4 HOUSEHOLD_CALCULATION_TTL_SECONDS = 86_400 diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py index 8a7ed5dc8..0076c8633 100644 --- a/policyengine_api/runtime_cache/reform_impacts.py +++ b/policyengine_api/runtime_cache/reform_impacts.py @@ -1,10 +1,10 @@ """Recoverable reform-impact caching and submission coordination.""" -from dataclasses import asdict, dataclass, replace -from datetime import datetime, timezone import hashlib import re import time +from dataclasses import asdict, dataclass, replace +from datetime import UTC, datetime from typing import Any from policyengine_api.observability.identifiers import normalize_observability_id @@ -81,8 +81,8 @@ def _datetime_to_wire(value: datetime | None) -> str | None: if value is None: return None if value.tzinfo is None: - value = value.replace(tzinfo=timezone.utc) - return value.astimezone(timezone.utc).isoformat() + value = value.replace(tzinfo=UTC) + return value.astimezone(UTC).isoformat() def _datetime_from_wire(value: Any) -> datetime | None: @@ -92,7 +92,7 @@ def _datetime_from_wire(value: Any) -> datetime | None: raise ValueError("invalid cached datetime") parsed = datetime.fromisoformat(value) if parsed.tzinfo is not None: - parsed = parsed.astimezone(timezone.utc).replace(tzinfo=None) + parsed = parsed.astimezone(UTC).replace(tzinfo=None) return parsed @@ -391,7 +391,7 @@ def _score(impact: CachedReformImpact) -> float: return time.time() value = impact.start_time if value.tzinfo is None: - value = value.replace(tzinfo=timezone.utc) + value = value.replace(tzinfo=UTC) return value.timestamp() def set( diff --git a/policyengine_api/runtime_cache/settings.py b/policyengine_api/runtime_cache/settings.py index 31131f5fa..35aba6b7d 100644 --- a/policyengine_api/runtime_cache/settings.py +++ b/policyengine_api/runtime_cache/settings.py @@ -2,15 +2,14 @@ from __future__ import annotations -from collections.abc import Mapping -from dataclasses import dataclass, field import os import re +from collections.abc import Mapping +from dataclasses import dataclass, field from urllib.parse import SplitResult, urlsplit from pydantic import SecretStr - RUNTIME_CACHE_MODE = "RUNTIME_CACHE_MODE" RUNTIME_CACHE_URL = "RUNTIME_CACHE_URL" RUNTIME_CACHE_CA_CERT = "RUNTIME_CACHE_CA_CERT" diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py index a50f909f8..a8be578f5 100644 --- a/policyengine_api/services/economy_comparison.py +++ b/policyengine_api/services/economy_comparison.py @@ -150,11 +150,10 @@ def decile_impact(baseline: dict, reform: dict) -> dict: ) rel_decile_dict = rel_income_change_by_decile.to_dict() avg_decile_dict = avg_income_change_by_decile.to_dict() - result = dict( + return dict( relative={int(k): v for k, v in rel_decile_dict.items()}, average={int(k): v for k, v in avg_decile_dict.items()}, ) - return result def wealth_decile_impact(baseline: dict, reform: dict) -> dict: @@ -191,11 +190,10 @@ def wealth_decile_impact(baseline: dict, reform: dict) -> dict: ) rel_decile_dict = rel_income_change_by_decile.to_dict() avg_decile_dict = avg_income_change_by_decile.to_dict() - result = dict( + return dict( relative={int(k): v for k, v in rel_decile_dict.items()}, average={int(k): v for k, v in avg_decile_dict.items()}, ) - return result def inequality_impact(baseline: dict, reform: dict) -> dict: @@ -610,11 +608,11 @@ def uk_constituency_breakdown( if selected_country is not None: if selected_country == "ENGLAND" and not code.startswith("E"): continue - elif selected_country == "SCOTLAND" and not code.startswith("S"): + if selected_country == "SCOTLAND" and not code.startswith("S"): continue - elif selected_country == "WALES" and not code.startswith("W"): + if selected_country == "WALES" and not code.startswith("W"): continue - elif selected_country == "NORTHERN_IRELAND" and not code.startswith("N"): + if selected_country == "NORTHERN_IRELAND" and not code.startswith("N"): continue weight: np.ndarray = weights[i] @@ -729,11 +727,11 @@ def uk_local_authority_breakdown( if selected_country is not None: if selected_country == "england" and not code.startswith("E"): continue - elif selected_country == "scotland" and not code.startswith("S"): + if selected_country == "scotland" and not code.startswith("S"): continue - elif selected_country == "wales" and not code.startswith("W"): + if selected_country == "wales" and not code.startswith("W"): continue - elif selected_country == "northern_ireland" and not code.startswith("N"): + if selected_country == "northern_ireland" and not code.startswith("N"): continue weight: np.ndarray = weights[i] @@ -847,7 +845,7 @@ def compare_economic_outputs( constituency_impact=constituency_impact_data, local_authority_impact=local_authority_impact_data, ) - elif baseline.get("type") == "cliff": + if baseline.get("type") == "cliff": return dict( baseline=dict( cliff_gap=baseline["cliff_gap"], @@ -858,3 +856,4 @@ def compare_economic_outputs( cliff_share=reform["cliff_share"], ), ) + return None diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index ce7112851..c6169f4cf 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -3,10 +3,12 @@ import json import uuid from enum import Enum -from typing import Any, Literal, Optional +from typing import Any, Literal import httpx from dotenv import load_dotenv +from pydantic import BaseModel, Field + from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, EXECUTION_STATUSES_FAILURE, @@ -20,12 +22,14 @@ get_valid_state_codes, normalize_us_region, ) -from policyengine_api.data.v1_models import ReformImpact from policyengine_api.data.places import validate_place_code +from policyengine_api.data.v1_models import ReformImpact from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint from policyengine_api.observability import ( runtime as observability_runtime, +) +from policyengine_api.observability import ( set_runtime_context, ) from policyengine_api.observability.stages import ( @@ -46,10 +50,9 @@ ReformImpactHandoffError, ReformImpactsService, ) -from policyengine_api.utils import budget_window as budget_window_utils -from policyengine_api.worker_spm import validate_worker_spm, validate_worker_result from policyengine_api.spm import SPMSelection, SPMValidationError -from pydantic import BaseModel, Field +from policyengine_api.utils import budget_window as budget_window_utils +from policyengine_api.worker_spm import validate_worker_result, validate_worker_spm load_dotenv() @@ -126,8 +129,8 @@ class EconomicImpactResult(BaseModel): """ status: ImpactStatus - data: Optional[dict] = None - message: Optional[str] = None + data: dict | None = None + message: str | None = None model_config = {"frozen": True} # Make model immutable @@ -170,14 +173,14 @@ class BudgetWindowEconomicImpactResult(BaseModel): """ status: ImpactStatus - data: Optional[dict] = None - progress: Optional[int] = None + data: dict | None = None + progress: int | None = None completed_years: list[str] = Field(default_factory=list) computing_years: list[str] = Field(default_factory=list) queued_years: list[str] = Field(default_factory=list) - message: Optional[str] = None - error: Optional[str] = None - cache_status: Optional[str] = None + message: str | None = None + error: str | None = None + cache_status: str | None = None model_config = {"frozen": True} @@ -195,7 +198,7 @@ def to_dict(self) -> dict[str, Any]: @classmethod def completed( - cls, data: dict, *, cache_status: Optional[str] = None + cls, data: dict, *, cache_status: str | None = None ) -> "BudgetWindowEconomicImpactResult": return cls( status=ImpactStatus.OK, @@ -213,7 +216,7 @@ def computing( computing_years: list[str], queued_years: list[str], message: str, - cache_status: Optional[str] = None, + cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": return cls( status=ImpactStatus.COMPUTING, @@ -231,10 +234,10 @@ def failed( cls, message: str, *, - completed_years: Optional[list[str]] = None, - computing_years: Optional[list[str]] = None, - queued_years: Optional[list[str]] = None, - cache_status: Optional[str] = None, + completed_years: list[str] | None = None, + computing_years: list[str] | None = None, + queued_years: list[str] | None = None, + cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": logger.log_struct({"message": message}, severity="ERROR") return cls( @@ -722,7 +725,7 @@ def _get_budget_window_result_from_batch_job_id( total_years: int, queued_years_on_submit: list[str], spm: dict | None = None, - cache_status: Optional[str] = None, + cache_status: str | None = None, observability_id: str | None, ) -> BudgetWindowEconomicImpactResult: resolved_observability_id = observability_id @@ -807,8 +810,8 @@ def _build_budget_window_computing_result( completed_years: list[str], computing_years: list[str], queued_years: list[str], - progress: Optional[int] = None, - cache_status: Optional[str] = None, + progress: int | None = None, + cache_status: str | None = None, ) -> BudgetWindowEconomicImpactResult: resolved_progress = progress if resolved_progress is None: @@ -1151,21 +1154,17 @@ def _get_previous_impacts( Fetch any previous simulation runs for the given policy reform. """ - previous_impacts: list[Any] = [] - previous_impacts = ( - self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( - country_id, - policy_id, - baseline_policy_id, - region, - dataset, - time_period, - options_hash, - self._build_options_hash_lookup_pattern(options_hash), - api_version, - ) + return self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( + country_id, + policy_id, + baseline_policy_id, + region, + dataset, + time_period, + options_hash, + self._build_options_hash_lookup_pattern(options_hash), + api_version, ) - return previous_impacts def _get_most_recent_impact( self, @@ -1205,12 +1204,11 @@ def _determine_impact_action( status = most_recent_impact.status if status == ImpactStatus.OK.value: return ImpactAction.COMPLETED - elif status == ImpactStatus.ERROR.value: + if status == ImpactStatus.ERROR.value: return ImpactAction.FAILED - elif status == ImpactStatus.COMPUTING.value: + if status == ImpactStatus.COMPUTING.value: return ImpactAction.COMPUTING - else: - raise ValueError(f"Unknown impact status: {status}") + raise ValueError(f"Unknown impact status: {status}") @observability_runtime.span( ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_HANDLE_EXECUTION_STATE) @@ -1220,7 +1218,7 @@ def _handle_execution_state( setup_options: EconomicImpactSetupOptions, execution_state: str, reform_impact: ReformImpact, - execution: Optional[Any] = None, + execution: Any | None = None, ) -> EconomicImpactResult: """ Handle the state of the execution and return the appropriate status and result. @@ -1250,7 +1248,7 @@ def _handle_execution_state( ) return EconomicImpactResult.completed(data=result) - elif execution_state in EXECUTION_STATUSES_FAILURE: + if execution_state in EXECUTION_STATUSES_FAILURE: # For Modal, try to get error message from execution error_message = "Simulation entrypoint execution failed" if ( @@ -1273,15 +1271,14 @@ def _handle_execution_state( ) return EconomicImpactResult.error(message=error_message) - elif execution_state in EXECUTION_STATUSES_PENDING: + if execution_state in EXECUTION_STATUSES_PENDING: logger.log_struct( {"message": "Sim API execution is still running"}, severity="INFO", ) return EconomicImpactResult.computing() - else: - raise ValueError(f"Unexpected sim API execution state: {execution_state}") + raise ValueError(f"Unexpected sim API execution state: {execution_state}") @observability_runtime.span( ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_READ_COMPLETED) @@ -1593,7 +1590,7 @@ def _with_policyengine_bundle( self, result: dict, setup_options: EconomicImpactSetupOptions, - execution: Optional[Any] = None, + execution: Any | None = None, ) -> dict: result = result if isinstance(result, dict) else {} cached_resolved_app_name = result.get("resolved_app_name") diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index bce5743c7..17ad32244 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -1,36 +1,37 @@ from __future__ import annotations +import time +from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass, replace from datetime import date -import time -from typing import Any, Callable +from typing import Any from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, POLICYENGINE_VERSION -from policyengine_api.observability import runtime as observability_runtime -from policyengine_api.observability.stages import HOUSEHOLD_STAGES, Stage from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ( Household, Policy, ) -from policyengine_api.runtime_cache.dependencies import get_runtime_cache_context +from policyengine_api.observability import runtime as observability_runtime +from policyengine_api.observability.stages import HOUSEHOLD_STAGES, Stage from policyengine_api.runtime_cache.core import record_cache_event +from policyengine_api.runtime_cache.dependencies import get_runtime_cache_context from policyengine_api.runtime_cache.household_calculations import ( CachedHouseholdCalculation, HouseholdCalculationCache, HouseholdCalculationIdentity, ) -from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs -from policyengine_api.utils.input_validation import find_unrecognized_inputs from policyengine_api.spm import ( normalize_spm_selection, resolved_spm_configuration, validate_spm_calculation_provenance, ) +from policyengine_api.utils.deprecated_inputs import drop_deprecated_inputs +from policyengine_api.utils.input_validation import find_unrecognized_inputs @dataclass(frozen=True) diff --git a/policyengine_api/services/household_mirroring.py b/policyengine_api/services/household_mirroring.py index 27760fb3d..863b89493 100644 --- a/policyengine_api/services/household_mirroring.py +++ b/policyengine_api/services/household_mirroring.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable import time +from collections.abc import Callable from typing import Protocol from sqlalchemy.exc import SQLAlchemyError diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index cc2adc33c..c9e35910e 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,10 +1,10 @@ from __future__ import annotations -from collections.abc import Callable import copy +from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import UTC, datetime from enum import StrEnum from sqlalchemy import select @@ -12,6 +12,8 @@ IntegrityError, OperationalError, SQLAlchemyError, +) +from sqlalchemy.exc import ( TimeoutError as SQLAlchemyTimeoutError, ) from sqlalchemy.orm import Session, sessionmaker @@ -26,8 +28,8 @@ LegacyHouseholdPersistenceResult, LegacyHouseholdSnapshot, ) -from policyengine_api.utils import hash_object from policyengine_api.spm import SPMValidationError, normalize_spm_selection +from policyengine_api.utils import hash_object def household_storage_json(country_id: str, household_json: dict, spm=None) -> dict: @@ -80,7 +82,7 @@ def retryable(self) -> bool: } @classmethod - def from_exception(cls, error: Exception) -> "HouseholdPersistenceError": + def from_exception(cls, error: Exception) -> HouseholdPersistenceError: if isinstance(error, SQLAlchemyTimeoutError): return cls(HouseholdPersistenceFailureCategory.TIMEOUT) if isinstance(error, OperationalError): @@ -289,7 +291,7 @@ def process_mirror_event( ) result = processor(pending) if event.processed_at is None: - event.processed_at = datetime.now(timezone.utc).replace(tzinfo=None) + event.processed_at = datetime.now(UTC).replace(tzinfo=None) session.add(event) session.flush() if after_processed_commit is not None: diff --git a/policyengine_api/services/policy_mirroring.py b/policyengine_api/services/policy_mirroring.py index 8963e001c..326312fb4 100644 --- a/policyengine_api/services/policy_mirroring.py +++ b/policyengine_api/services/policy_mirroring.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable import time +from collections.abc import Callable from typing import Protocol from sqlalchemy.exc import SQLAlchemyError @@ -12,6 +12,8 @@ MetadataCatalogUnavailableError, MetadataCatalogVersionNotFoundError, ) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.gcp_logging import logger from policyengine_api.services.v2.policies.types import ( LegacyPolicyPersistenceResult, LegacyPolicySnapshot, @@ -23,8 +25,6 @@ PolicyContentHashCollisionError, PolicyCreationIntegrityError, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.gcp_logging import logger class LegacyPolicyMirror(Protocol): diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index e7af296e3..9a2d2ab5d 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,8 +1,8 @@ from __future__ import annotations import copy -from dataclasses import dataclass from collections.abc import Iterator +from dataclasses import dataclass from typing import Any from sqlalchemy import select diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 510895636..a046f060e 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -294,4 +294,4 @@ def set_complete_reform_impact( @staticmethod def _now() -> datetime.datetime: - return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + return datetime.datetime.now(datetime.UTC).replace(tzinfo=None) diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 69246e827..b218bf39b 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,6 +1,6 @@ import json from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import UTC, datetime from typing import cast from sqlalchemy import select @@ -116,7 +116,7 @@ def _sessions(self) -> sessionmaker[Session]: @staticmethod def _utc_timestamp() -> datetime: - return datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) + return datetime.now(UTC).replace(microsecond=0, tzinfo=None) @staticmethod def format_run_timestamp(value: datetime | str | None) -> str | None: @@ -125,9 +125,9 @@ def format_run_timestamp(value: datetime | str | None) -> str | None: if isinstance(value, datetime): timestamp = value if timestamp.tzinfo is None: - timestamp = timestamp.replace(tzinfo=timezone.utc) + timestamp = timestamp.replace(tzinfo=UTC) return ( - timestamp.astimezone(timezone.utc) + timestamp.astimezone(UTC) .replace(microsecond=0) .isoformat() .replace("+00:00", "Z") @@ -144,9 +144,9 @@ def format_run_timestamp(value: datetime | str | None) -> str | None: except ValueError: return normalized if normalized.endswith("Z") else f"{normalized}Z" if parsed.tzinfo is None: - parsed = parsed.replace(tzinfo=timezone.utc) + parsed = parsed.replace(tzinfo=UTC) return ( - parsed.astimezone(timezone.utc) + parsed.astimezone(UTC) .replace(microsecond=0) .isoformat() .replace("+00:00", "Z") diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index e0c8127ff..aaf6daf53 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -1,5 +1,5 @@ import uuid -from datetime import datetime, timezone +from datetime import UTC, datetime from typing import Any from sqlalchemy import func, select @@ -7,7 +7,6 @@ from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun - REPORT_RUN_VERSION_FIELDS = ( "country_package_version", "policyengine_version", @@ -60,7 +59,7 @@ def create_report_output_run( ) or 0 ) + 1 - now = datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None) + now = datetime.now(UTC).replace(microsecond=0, tzinfo=None) values = { "status": status, "output": output, diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index 93790586f..6bcb82609 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -7,7 +7,6 @@ from policyengine_api.data.v1_models import ReportOutput, Simulation from policyengine_api.utils.population_identity import canonical_numeric_household_id - REPORT_SPEC_SCHEMA_VERSION = 1 REPORT_SPEC_STATUSES = {"explicit", "backfilled_assumed"} HOUSEHOLD_REPORT_KINDS = {"household_single", "household_comparison"} diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 61b3a970b..677f2b6d5 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -6,7 +6,6 @@ from policyengine_api.data.v1_models import Simulation, SimulationRun - SIMULATION_RUN_VERSION_FIELDS = ( "country_package_version", "policyengine_version", diff --git a/policyengine_api/services/user_policy_mirroring.py b/policyengine_api/services/user_policy_mirroring.py index 7e0e3c3d9..1e44e1a06 100644 --- a/policyengine_api/services/user_policy_mirroring.py +++ b/policyengine_api/services/user_policy_mirroring.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable import time +from collections.abc import Callable from typing import Protocol from sqlalchemy.exc import SQLAlchemyError @@ -12,6 +12,13 @@ MetadataCatalogUnavailableError, MetadataCatalogVersionNotFoundError, ) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.gcp_logging import logger +from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.user_policy_service import ( + PendingUserPolicyMirrorEvent, + UserPolicyService, +) from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot from policyengine_api.services.v2.policies.validators import ( LegacyPolicyMappingIntegrityError, @@ -27,13 +34,6 @@ from policyengine_api.services.v2.user_policies.validators import ( LegacyUserPolicyIntegrityError, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.gcp_logging import logger -from policyengine_api.services.policy_service import PolicyService -from policyengine_api.services.user_policy_service import ( - PendingUserPolicyMirrorEvent, - UserPolicyService, -) class LegacyUserPolicyMirror(Protocol): diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py index ab0c4e7d3..42d60150a 100644 --- a/policyengine_api/services/user_policy_service.py +++ b/policyengine_api/services/user_policy_service.py @@ -1,30 +1,31 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import UTC, datetime from enum import StrEnum -from typing import Any, Mapping +from typing import Any from sqlalchemy import select from sqlalchemy.exc import ( IntegrityError, OperationalError, SQLAlchemyError, +) +from sqlalchemy.exc import ( TimeoutError as SQLAlchemyTimeoutError, ) from sqlalchemy.orm import Session, sessionmaker from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import UserPolicy, UserPolicyMirrorEvent +from policyengine_api.services.v2.user_policies.transformations import ( + fingerprint_legacy_user_policy, +) from policyengine_api.services.v2.user_policies.types import ( LegacyUserPolicyPersistenceResult, LegacyUserPolicySnapshot, ) -from policyengine_api.services.v2.user_policies.transformations import ( - fingerprint_legacy_user_policy, -) - USER_POLICY_IDENTITY_FIELDS = ( "country_id", @@ -406,7 +407,7 @@ def process_pending_mirror_events( break pending = self._decode_mirror_event(event) processed_result = processor(pending) - event.processed_at = datetime.now(timezone.utc).replace(tzinfo=None) + event.processed_at = datetime.now(UTC).replace(tzinfo=None) session.add(event) session.flush() latest_result = processed_result diff --git a/policyengine_api/services/v2/comparison_runs/__init__.py b/policyengine_api/services/v2/comparison_runs/__init__.py index 166221914..a0e8f6084 100644 --- a/policyengine_api/services/v2/comparison_runs/__init__.py +++ b/policyengine_api/services/v2/comparison_runs/__init__.py @@ -1,10 +1,10 @@ """Temporary Stage 12 comparison-run lifecycle contracts.""" from policyengine_api.services.v2.comparison_runs.types import ( - ComparisonRunAggregationStatus, - ComparisonRunLifecycleStatus, ComparisonReportPersistenceResult, ComparisonReportRecord, + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, ComparisonSimulationPersistenceResult, ComparisonSimulationRecord, ResultComparisonStatus, diff --git a/policyengine_api/services/v2/comparison_runs/services.py b/policyengine_api/services/v2/comparison_runs/services.py index 9f9737cb2..1437d2f41 100644 --- a/policyengine_api/services/v2/comparison_runs/services.py +++ b/policyengine_api/services/v2/comparison_runs/services.py @@ -29,9 +29,9 @@ simulation_record, ) from policyengine_api.services.v2.comparison_runs.types import ( - ComparisonRunLifecycleStatus, ComparisonReportPersistenceResult, ComparisonReportRecord, + ComparisonRunLifecycleStatus, ComparisonSimulationPersistenceResult, ComparisonSimulationRecord, ) diff --git a/policyengine_api/services/v2/comparison_runs/validators.py b/policyengine_api/services/v2/comparison_runs/validators.py index 7563bb62a..756a8b41b 100644 --- a/policyengine_api/services/v2/comparison_runs/validators.py +++ b/policyengine_api/services/v2/comparison_runs/validators.py @@ -5,9 +5,9 @@ from collections.abc import Iterable from policyengine_api.services.v2.comparison_runs.types import ( + ComparisonReportRecord, ComparisonRunAggregationStatus, ComparisonRunLifecycleStatus, - ComparisonReportRecord, ComparisonSimulationRecord, ResultComparisonStatus, ) diff --git a/policyengine_api/services/v2/households/transformations.py b/policyengine_api/services/v2/households/transformations.py index f969ad160..43c15b06b 100644 --- a/policyengine_api/services/v2/households/transformations.py +++ b/policyengine_api/services/v2/households/transformations.py @@ -2,12 +2,14 @@ from __future__ import annotations -from decimal import Decimal import hashlib import json +from decimal import Decimal from typing import Any, cast -from policyengine_core.periods import period as parse_policyengine_period # type: ignore[import-untyped] +from policyengine_core.periods import ( # type: ignore[import-untyped] + period as parse_policyengine_period, +) from policyengine_api.data.v2.models import Household from policyengine_api.services.v2.households.types import ( @@ -18,12 +20,11 @@ LegacyHouseholdSnapshot, ) from policyengine_api.services.v2.households.validators import ( - LegacyHouseholdTranslationError, SUPPORTED_ENTITY_COLLECTIONS, + LegacyHouseholdTranslationError, normalize_household_document, ) - HOUSEHOLD_CANONICALIZATION_VERSION = 1 LEGACY_HOUSEHOLD_FINGERPRINT_VERSION = 1 LEGACY_COLLECTION_NAMES = { @@ -36,7 +37,7 @@ } -def _canonical_number(value: int | float) -> str: +def _canonical_number(value: float) -> str: number = Decimal(str(value)) if number.is_zero(): return "0" diff --git a/policyengine_api/services/v2/households/types.py b/policyengine_api/services/v2/households/types.py index 40ca5a30d..17b50d609 100644 --- a/policyengine_api/services/v2/households/types.py +++ b/policyengine_api/services/v2/households/types.py @@ -19,7 +19,6 @@ from policyengine_api.query_parameters import CountryId, DefaultYear from policyengine_api.services.v2.households.validators import require_json_value - StrictJsonValue = Annotated[JsonValue, BeforeValidator(require_json_value)] diff --git a/policyengine_api/services/v2/households/validators.py b/policyengine_api/services/v2/households/validators.py index 7c6040c7a..ee0abefc6 100644 --- a/policyengine_api/services/v2/households/validators.py +++ b/policyengine_api/services/v2/households/validators.py @@ -9,7 +9,6 @@ from policyengine_api.data.v2.models import LegacyHouseholdMapping from policyengine_api.spm import SPMSelection - MAXIMUM_HOUSEHOLD_ENTITIES = 1_000 MAXIMUM_HOUSEHOLD_RECORD_VALUES = 500 MAXIMUM_JSON_NESTING = 100 diff --git a/policyengine_api/services/v2/metadata/services.py b/policyengine_api/services/v2/metadata/services.py index 1e24861f5..a0f9e625a 100644 --- a/policyengine_api/services/v2/metadata/services.py +++ b/policyengine_api/services/v2/metadata/services.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from uuid import UUID from policyengine_api.data.v2.catalog.catalog_selection import ( @@ -312,9 +312,7 @@ def list_parameter_values( selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) - selected_day = ( - utc_day_start(now or datetime.now(timezone.utc)) if current else None - ) + selected_day = utc_day_start(now or datetime.now(UTC)) if current else None rows = read_parameter_values( self._database_session.session, model_version_id=selected.model_version.id, diff --git a/policyengine_api/services/v2/metadata/transformations.py b/policyengine_api/services/v2/metadata/transformations.py index 83208f41f..e71ef90f4 100644 --- a/policyengine_api/services/v2/metadata/transformations.py +++ b/policyengine_api/services/v2/metadata/transformations.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from typing import Any, TypeVar from policyengine_api.data.v2.catalog.catalog_selection import SelectedCatalog @@ -32,7 +32,6 @@ MetadataVariable, ) - ResourceT = TypeVar("ResourceT") @@ -117,8 +116,8 @@ def metadata_parameter_value(value: ParameterValue) -> MetadataCanonicalParamete def utc_day_start(selected_time: datetime) -> datetime: if selected_time.tzinfo is None: - selected_time = selected_time.replace(tzinfo=timezone.utc) - return selected_time.astimezone(timezone.utc).replace( + selected_time = selected_time.replace(tzinfo=UTC) + return selected_time.astimezone(UTC).replace( hour=0, minute=0, second=0, microsecond=0 ) diff --git a/policyengine_api/services/v2/metadata/validators.py b/policyengine_api/services/v2/metadata/validators.py index 1ca60c3a7..6172d2018 100644 --- a/policyengine_api/services/v2/metadata/validators.py +++ b/policyengine_api/services/v2/metadata/validators.py @@ -41,12 +41,12 @@ def require_metadata_resource( def validate_economy_options( - selected: "SelectedCatalog", + selected: SelectedCatalog, *, country_id: str, - regions: list["Region"], - national_dataset: "Dataset | None", -) -> "Dataset": + regions: list[Region], + national_dataset: Dataset | None, +) -> Dataset: """Validate preloaded records required by the economy-options response.""" if not any(region.code == country_id for region in regions): diff --git a/policyengine_api/services/v2/policies/transformations.py b/policyengine_api/services/v2/policies/transformations.py index 576ae30c8..1e4a93488 100644 --- a/policyengine_api/services/v2/policies/transformations.py +++ b/policyengine_api/services/v2/policies/transformations.py @@ -2,15 +2,17 @@ from __future__ import annotations -from collections.abc import Mapping -from datetime import date, datetime, time, timezone -from decimal import Decimal import hashlib import json +from collections.abc import Mapping +from datetime import UTC, date, datetime, time +from decimal import Decimal from typing import Any from uuid import UUID -from policyengine_core.periods import period as parse_policyengine_period # type: ignore[import-untyped] +from policyengine_core.periods import ( # type: ignore[import-untyped] + period as parse_policyengine_period, +) from policyengine_api.data.v2.catalog.catalog_selection import SelectedCatalog from policyengine_api.data.v2.models import ( @@ -34,11 +36,10 @@ validate_policy_catalog, ) - POLICY_CANONICALIZATION_VERSION = 1 -def _canonical_number(value: int | float) -> str: +def _canonical_number(value: float) -> str: number = Decimal(str(value)) if number.is_zero(): return "0" @@ -69,7 +70,7 @@ def _canonical_json(value: Any) -> str: def canonical_utc_datetime(value: datetime) -> str: - utc_value = value.astimezone(timezone.utc) + utc_value = value.astimezone(UTC) return utc_value.isoformat(timespec="microseconds").replace("+00:00", "Z") @@ -200,7 +201,7 @@ def _utc_midnight(value: str) -> datetime: raise LegacyPolicyTranslationError( f"legacy period {value!r} is invalid" ) from error - return datetime.combine(parsed, time.min, tzinfo=timezone.utc) + return datetime.combine(parsed, time.min, tzinfo=UTC) def parse_legacy_period(value: str) -> tuple[datetime, datetime]: diff --git a/policyengine_api/services/v2/policies/types.py b/policyengine_api/services/v2/policies/types.py index 86dbc19bd..ef1f2823a 100644 --- a/policyengine_api/services/v2/policies/types.py +++ b/policyengine_api/services/v2/policies/types.py @@ -25,7 +25,6 @@ require_json_value, ) - StrictJsonValue = Annotated[JsonValue, BeforeValidator(require_json_value)] UtcDateTime = Annotated[datetime, AfterValidator(normalize_utc)] @@ -45,7 +44,7 @@ class PolicyParameterValueInput(StrictPolicyInput): end_date: UtcDateTime | None = None @model_validator(mode="after") - def validate_period(self) -> "PolicyParameterValueInput": + def validate_period(self) -> PolicyParameterValueInput: if self.end_date is not None and self.end_date < self.start_date: raise ValueError("end_date must not precede start_date") return self diff --git a/policyengine_api/services/v2/policies/validators.py b/policyengine_api/services/v2/policies/validators.py index 822db40b5..ac0322e82 100644 --- a/policyengine_api/services/v2/policies/validators.py +++ b/policyengine_api/services/v2/policies/validators.py @@ -2,8 +2,8 @@ from __future__ import annotations -from datetime import datetime, timezone import math +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any from uuid import UUID @@ -83,15 +83,15 @@ def normalize_utc(value: datetime) -> datetime: if value.tzinfo is None or value.utcoffset() is None: raise ValueError("effective dates must include a UTC offset") - return value.astimezone(timezone.utc) + return value.astimezone(UTC) def validate_policy_catalog( - policy_input: "PolicyCreationInput", + policy_input: PolicyCreationInput, *, selected: SelectedCatalog, resolved_parameter_ids: set[UUID], -) -> "ResolvedPolicyCreationInput": +) -> ResolvedPolicyCreationInput: """Validate preloaded catalog records and bind them to policy content.""" from policyengine_api.services.v2.policies.types import ResolvedPolicyCreationInput diff --git a/policyengine_api/services/v2/user_households/types.py b/policyengine_api/services/v2/user_households/types.py index 3765dc1ac..095be46eb 100644 --- a/policyengine_api/services/v2/user_households/types.py +++ b/policyengine_api/services/v2/user_households/types.py @@ -11,7 +11,6 @@ from policyengine_api.query_parameters import CountryId, ResourceId, UserId - AssociationName = Annotated[str, StringConstraints(max_length=255)] @@ -33,7 +32,7 @@ class UserHouseholdUpdateInput(StrictUserHouseholdInput): description: str | None = None @model_validator(mode="after") - def require_valid_supplied_field(self) -> "UserHouseholdUpdateInput": + def require_valid_supplied_field(self) -> UserHouseholdUpdateInput: mutable_fields = {"household_id", "name", "description"} supplied = self.model_fields_set.intersection(mutable_fields) if not supplied: diff --git a/policyengine_api/services/v2/user_policies/transformations.py b/policyengine_api/services/v2/user_policies/transformations.py index 615059f07..230f7a27e 100644 --- a/policyengine_api/services/v2/user_policies/transformations.py +++ b/policyengine_api/services/v2/user_policies/transformations.py @@ -2,8 +2,8 @@ from __future__ import annotations -from hashlib import sha256 import json +from hashlib import sha256 from uuid import UUID from policyengine_api.data.v2.models import UserPolicy @@ -14,7 +14,6 @@ UserPolicyRead, ) - USER_POLICY_FINGERPRINT_VERSION = 1 diff --git a/policyengine_api/services/v2/user_policies/types.py b/policyengine_api/services/v2/user_policies/types.py index 67f70cc6a..5a4a97f94 100644 --- a/policyengine_api/services/v2/user_policies/types.py +++ b/policyengine_api/services/v2/user_policies/types.py @@ -37,7 +37,7 @@ class UserPolicyUpdateInput(StrictAssociationInput): description: str | None = None @model_validator(mode="after") - def require_supplied_field(self) -> "UserPolicyUpdateInput": + def require_supplied_field(self) -> UserPolicyUpdateInput: if not self.model_fields_set.intersection({"name", "description"}): raise ValueError("At least one of name or description must be supplied") return self diff --git a/policyengine_api/services/v2/user_policies/validators.py b/policyengine_api/services/v2/user_policies/validators.py index 1d895ecba..3172750e8 100644 --- a/policyengine_api/services/v2/user_policies/validators.py +++ b/policyengine_api/services/v2/user_policies/validators.py @@ -55,7 +55,7 @@ def require_user_policy( def validate_association_creation( - association_input: "UserPolicyCreationInput", + association_input: UserPolicyCreationInput, *, user: User | None, policy: Policy | None, @@ -95,8 +95,8 @@ def validate_legacy_user_mapping_conflict( def validate_legacy_user_policy_input( - snapshot: "LegacyUserPolicySnapshot", - reform_snapshot: "LegacyPolicySnapshot", + snapshot: LegacyUserPolicySnapshot, + reform_snapshot: LegacyPolicySnapshot, *, source_revision: int, ) -> None: diff --git a/policyengine_api/specification.py b/policyengine_api/specification.py index e2c9296b6..1d4b378a5 100644 --- a/policyengine_api/specification.py +++ b/policyengine_api/specification.py @@ -7,6 +7,7 @@ from typing import cast import yaml + from policyengine_api.constants import VERSION from policyengine_api.json_types import JSONObject from policyengine_api.query_parameters import ( @@ -21,7 +22,6 @@ SPMSelection, ) - DEFAULT_SPECIFICATION_PATH = Path(__file__).with_name("openapi_spec.yaml") diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py index 4f92a3b96..0130931ed 100644 --- a/policyengine_api/spm.py +++ b/policyengine_api/spm.py @@ -116,7 +116,7 @@ def validate_as_of(cls, value: str | None) -> str | None: return value @model_validator(mode="after") - def validate_location(self) -> "SPMResolvedConfiguration": + def validate_location(self) -> SPMResolvedConfiguration: if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -138,7 +138,7 @@ class SPMProvenance(BaseModel): county_vintage: str = Field(pattern=r"^[0-9]{4}$") as_of: date | None years: tuple[str, ...] - runtime_versions: "SPMRuntimeVersions" + runtime_versions: SPMRuntimeVersions @field_validator("years") @classmethod @@ -152,7 +152,7 @@ def validate_years(cls, years: tuple[str, ...]) -> tuple[str, ...]: return years @model_validator(mode="after") - def validate_location(self) -> "SPMProvenance": + def validate_location(self) -> SPMProvenance: if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("A metro SPM receipt requires geography_id") @@ -191,7 +191,7 @@ class SPMComparisonProvenance(BaseModel): reform: SPMExecutionProvenance @model_validator(mode="after") - def validate_compatible_receipts(self) -> "SPMComparisonProvenance": + def validate_compatible_receipts(self) -> SPMComparisonProvenance: if self.baseline.receipt != self.reform.receipt: raise ValueError("baseline and reform SPM receipts differ") return self diff --git a/policyengine_api/utils/__init__.py b/policyengine_api/utils/__init__.py index 58a8dbb62..1aeb79c92 100644 --- a/policyengine_api/utils/__init__.py +++ b/policyengine_api/utils/__init__.py @@ -1,4 +1,4 @@ -from .json import * # noqa: F403 from .cache_utils import * # noqa: F403 -from .singleton import Singleton as Singleton from .get_current_law import get_current_law_policy_id as get_current_law_policy_id +from .json import * # noqa: F403 +from .singleton import Singleton as Singleton diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py index 12f3c6d04..c6da3812a 100644 --- a/policyengine_api/utils/cache_utils.py +++ b/policyengine_api/utils/cache_utils.py @@ -3,6 +3,7 @@ import hashlib import json import logging + import flask diff --git a/policyengine_api/utils/deprecated_inputs.py b/policyengine_api/utils/deprecated_inputs.py index 79dbecdf7..2b45b4573 100644 --- a/policyengine_api/utils/deprecated_inputs.py +++ b/policyengine_api/utils/deprecated_inputs.py @@ -11,7 +11,6 @@ import copy from dataclasses import dataclass - # Registry of removed/renamed model variables that legacy partner traffic # may still pass. The value is a one-line migration hint surfaced verbatim # in the warning message - keep it actionable. diff --git a/policyengine_api/utils/json.py b/policyengine_api/utils/json.py index b69c5a1dc..294e797db 100644 --- a/policyengine_api/utils/json.py +++ b/policyengine_api/utils/json.py @@ -1,11 +1,12 @@ -import hashlib import base64 +import hashlib + import numpy as np def make_hashable(o): if isinstance(o, (tuple, list)): - return tuple((make_hashable(e) for e in o)) + return tuple(make_hashable(e) for e in o) if isinstance(o, dict): return tuple(sorted((k, make_hashable(v)) for k, v in o.items())) @@ -28,7 +29,7 @@ def get_safe_json(value): if isinstance(value, (int, float)): if value == np.inf: return "Infinity" - elif value == -np.inf: + if value == -np.inf: return "-Infinity" return value if isinstance(value, str): diff --git a/policyengine_api/utils/payload_validators/__init__.py b/policyengine_api/utils/payload_validators/__init__.py index 1c3418be7..6da6b53cb 100644 --- a/policyengine_api/utils/payload_validators/__init__.py +++ b/policyengine_api/utils/payload_validators/__init__.py @@ -1,7 +1,7 @@ from .validate_country import validate_country as validate_country -from .validate_set_policy_payload import ( - validate_set_policy_payload as validate_set_policy_payload, -) from .validate_household_payload import ( validate_household_payload as validate_household_payload, ) +from .validate_set_policy_payload import ( + validate_set_policy_payload as validate_set_policy_payload, +) diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index 346840d34..adc2204d0 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -1,5 +1,4 @@ from functools import wraps -from typing import Union from flask import Response @@ -21,9 +20,7 @@ def validate_country(func): """ @wraps(func) - def validate_country_wrapper( - country_id: str, *args, **kwargs - ) -> Union[None, Response]: + def validate_country_wrapper(country_id: str, *args, **kwargs) -> None | Response: print("Validating country") try: ensure_supported_country(country_id) diff --git a/policyengine_api/utils/singleton.py b/policyengine_api/utils/singleton.py index 3776cb92d..9767a5c7e 100644 --- a/policyengine_api/utils/singleton.py +++ b/policyengine_api/utils/singleton.py @@ -3,5 +3,5 @@ class Singleton(type): def __call__(cls, *args, **kwargs): if cls not in cls._instances: - cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs) + cls._instances[cls] = super().__call__(*args, **kwargs) return cls._instances[cls] diff --git a/policyengine_api/worker_spm_release.py b/policyengine_api/worker_spm_release.py index 8f876e0d2..79115b8c1 100644 --- a/policyengine_api/worker_spm_release.py +++ b/policyengine_api/worker_spm_release.py @@ -2,14 +2,13 @@ from __future__ import annotations -from importlib.metadata import version as distribution_version import json import sys +from importlib.metadata import version as distribution_version from policyengine_api import constants from policyengine_api.worker_spm import validate_worker_spm - REQUIRED_BUNDLE_PACKAGES = ( "policyengine", "policyengine-core", diff --git a/pyproject.toml b/pyproject.toml index 05b3912f5..7067ae5be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,12 +80,19 @@ packages = ["policyengine_api"] allow-direct-references = true [tool.ruff.lint] -# Ruff's defaults as of 0.15, made explicit so a ruff upgrade cannot change -# which rules apply. Widen this deliberately, one rule family per PR. -select = ["E4", "E7", "E9", "F"] +# Rule families are enabled one themed group per PR, each fixing everything it +# finds (no ratchet), until the select list is "ALL". See +# docs/engineering/skills/rules-registry.md and rules/registry.toml. +select = [ + "E4", "E7", "E9", "F", + # Group 1: safe auto-fixes. + "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W", +] [tool.ruff.lint.per-file-ignores] -"migrations/*/versions/*.py" = ["F401"] +# Alembic revisions must stay exactly as Alembic generates them +# (api.migration.alembic-autogenerated-only), so ruff does not lint them. +"migrations/*/versions/*.py" = ["ALL"] [tool.mypy] python_version = "3.11" diff --git a/rules/registry.toml b/rules/registry.toml index 9495ef0ee..e7a4c6bdb 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -670,21 +670,25 @@ invocation = "ruff format --check" [[rule]] id = "api.quality.ruff-lint" category = "quality" -statement = "Changed Python files pass ruff check with the rules selected in pyproject.toml (ruff's defaults, made explicit); agents run it before committing, but CI does not yet." -rationale = "docs/engineering/skills/github-prs.md#commit-hygiene" +statement = "All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL." +rationale = "docs/engineering/skills/rules-registry.md#baselines" scope = ["**/*.py"] tier = "mechanical" -status = "advisory" +status = "enforced" autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["E4", "E7", "E9", "F"] +codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W"] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "lint" +invocation = "ruff check ." [[rule]] id = "api.quality.alembic-revisions-unused-imports" category = "quality" -statement = "Alembic revision files are exempt from unused-import checks (F401), because autogenerated revisions import types they may not use." +statement = "Alembic revision files are exempt from ruff lint rules, because they must stay exactly as Alembic generates them." rationale = "docs/engineering/skills/alembic-migrations.md#mandatory-generation-rule" scope = ["migrations/*/versions/*.py"] tier = "mechanical" @@ -694,7 +698,7 @@ autofix = false kind = "ruff" mode = "per-file-ignore" pattern = "migrations/*/versions/*.py" -codes = ["F401"] +codes = ["ALL"] [[rule]] id = "api.quality.mypy-v2" diff --git a/scripts/capture_migration_baseline.py b/scripts/capture_migration_baseline.py index e7c7beebf..24662190b 100644 --- a/scripts/capture_migration_baseline.py +++ b/scripts/capture_migration_baseline.py @@ -12,9 +12,9 @@ import os import statistics import time +from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path -from typing import Iterable import httpx diff --git a/scripts/export_migration_contracts.py b/scripts/export_migration_contracts.py index 218fe3f4c..0c5e50711 100644 --- a/scripts/export_migration_contracts.py +++ b/scripts/export_migration_contracts.py @@ -9,7 +9,6 @@ from pathlib import Path from typing import Any - REPO_ROOT = Path(__file__).resolve().parent.parent if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -17,7 +16,6 @@ from policyengine_api.migration_registry import ROUTE_GROUPS # noqa: E402 from tests.contract.registry import APP_V2_WORKFLOW_CONTRACTS # noqa: E402 - DEFAULT_JSON = REPO_ROOT / "docs" / "generated" / "migration_contracts.json" DEFAULT_MARKDOWN = REPO_ROOT / "docs" / "engineering" / "migration-contracts.md" diff --git a/scripts/guards/migration_contracts.py b/scripts/guards/migration_contracts.py index 04d3c0918..2e61125c5 100644 --- a/scripts/guards/migration_contracts.py +++ b/scripts/guards/migration_contracts.py @@ -7,7 +7,6 @@ from scripts import export_migration_contracts - REPO_ROOT = Path(__file__).resolve().parents[2] ALLOWED_CURRENT_CONTRACTS = frozenset({"api_v1_compatible", "typed_v2_resources"}) diff --git a/scripts/measure_cloud_run_runtime.py b/scripts/measure_cloud_run_runtime.py index 56bfb8d4d..e269315b0 100644 --- a/scripts/measure_cloud_run_runtime.py +++ b/scripts/measure_cloud_run_runtime.py @@ -17,12 +17,12 @@ import asyncio import copy import json -from collections import Counter import random import sys import time import uuid -from datetime import datetime, timezone +from collections import Counter +from datetime import UTC, datetime from pathlib import Path import httpx @@ -47,7 +47,7 @@ def _utc_now_iso() -> str: - return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") def _classify(status_code: int) -> str: @@ -206,7 +206,7 @@ async def _run(args: argparse.Namespace) -> dict: ended_at = _utc_now_iso() measured = [r for r in records if not r["warmup"]] - summary = { + return { "cell": args.cell_label, "base_url": args.base_url, "clients": args.clients, @@ -230,7 +230,6 @@ async def _run(args: argparse.Namespace) -> dict: for endpoint in ("metadata", "calculate") }, } - return summary def main() -> int: diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index e9d897632..566144d49 100644 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -4,10 +4,10 @@ from __future__ import annotations import argparse -from collections.abc import Mapping, Sequence import json import os import sys +from collections.abc import Mapping, Sequence from sqlalchemy import Engine, create_engine from sqlalchemy.engine import URL @@ -15,10 +15,10 @@ from policyengine_api.data.orm import build_session_factory from policyengine_api.gcp_logging import logger -from policyengine_api.services.household_service import HouseholdService from policyengine_api.services.household_mirroring import ( process_household_event_after_commit, ) +from policyengine_api.services.household_service import HouseholdService class HouseholdEventCommandConfigurationError(RuntimeError): diff --git a/scripts/publish_v2_metadata_catalog.py b/scripts/publish_v2_metadata_catalog.py index aab5212d4..e9a2ffa4a 100644 --- a/scripts/publish_v2_metadata_catalog.py +++ b/scripts/publish_v2_metadata_catalog.py @@ -3,6 +3,5 @@ from policyengine_api.data.v2.catalog.initialization import main - if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/qualify_v2_household_migration.py b/scripts/qualify_v2_household_migration.py index f21cc46ea..1e99a15d8 100644 --- a/scripts/qualify_v2_household_migration.py +++ b/scripts/qualify_v2_household_migration.py @@ -3,6 +3,5 @@ from policyengine_api.data.v2.household_migration_qualification import main - if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/qualify_v2_policy_migration.py b/scripts/qualify_v2_policy_migration.py index 841ee20ea..d7a1743ec 100644 --- a/scripts/qualify_v2_policy_migration.py +++ b/scripts/qualify_v2_policy_migration.py @@ -3,6 +3,5 @@ from policyengine_api.data.v2.policy_migration_qualification import main - if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/rules/tests/support.py b/scripts/rules/tests/support.py index 42f9c257f..025683974 100644 --- a/scripts/rules/tests/support.py +++ b/scripts/rules/tests/support.py @@ -6,6 +6,7 @@ import textwrap from datetime import date from pathlib import Path +from typing import Self from scripts.rules.guard import GuardResult, run_guard from scripts.rules.load import load_registry @@ -195,7 +196,7 @@ def __init__(self, files: dict[str, str] | None = None, *, render: bool = True): if render: self.render() - def __enter__(self) -> Repo: + def __enter__(self) -> Self: return self def __exit__(self, *exc: object) -> None: diff --git a/scripts/run_quality_guards.py b/scripts/run_quality_guards.py index 464460000..5da7c5fb5 100644 --- a/scripts/run_quality_guards.py +++ b/scripts/run_quality_guards.py @@ -6,7 +6,6 @@ from collections.abc import Callable from pathlib import Path - REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -15,7 +14,6 @@ from scripts.rules import ci_scripts # noqa: E402 from scripts.rules import guard as rules_guard # noqa: E402 - Guard = tuple[str, Callable[[], list[str]]] GUARDS: tuple[Guard, ...] = ( diff --git a/scripts/smoke_test_local_boot.py b/scripts/smoke_test_local_boot.py index c57cf6d9b..5dc8fa2aa 100644 --- a/scripts/smoke_test_local_boot.py +++ b/scripts/smoke_test_local_boot.py @@ -10,7 +10,6 @@ import importlib.metadata import os -from pathlib import Path import socket import subprocess import sys @@ -18,7 +17,7 @@ import time import urllib.error import urllib.request - +from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] HEALTH_PATHS = ("/liveness-check", "/readiness-check") diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index 80611b9e8..084792311 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -3,8 +3,8 @@ from __future__ import annotations import argparse -from enum import StrEnum import os +from enum import StrEnum from typing import Any from alembic import command @@ -19,7 +19,6 @@ from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base - ALEMBIC_CONFIG = REPO / "alembic-v1.ini" MIGRATION_LOCK_NAME = "policyengine-api-v1-alembic" DATABASE_NAME = "policyengine" diff --git a/tests/conftest.py b/tests/conftest.py index a45e6afb3..5bf715308 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,6 @@ import os -from pathlib import Path import sys +from pathlib import Path from unittest.mock import patch import pytest diff --git a/tests/contract/clients.py b/tests/contract/clients.py index 4fdd98500..2472abd80 100644 --- a/tests/contract/clients.py +++ b/tests/contract/clients.py @@ -1,7 +1,8 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass -from typing import Mapping, Protocol +from typing import Protocol from fastapi.testclient import TestClient from flask import Flask diff --git a/tests/contract/conftest.py b/tests/contract/conftest.py index 45a720dc3..842122ca3 100644 --- a/tests/contract/conftest.py +++ b/tests/contract/conftest.py @@ -1,4 +1,3 @@ import os - os.environ.setdefault("FLASK_DEBUG", "1") diff --git a/tests/contract/helpers.py b/tests/contract/helpers.py index 285450e80..6a04d595c 100644 --- a/tests/contract/helpers.py +++ b/tests/contract/helpers.py @@ -3,7 +3,6 @@ import json from typing import Any - VOLATILE_STRING = "" diff --git a/tests/contract/test_policy_v2_compatibility.py b/tests/contract/test_policy_v2_compatibility.py index b8bfaec68..727b0dae6 100644 --- a/tests/contract/test_policy_v2_compatibility.py +++ b/tests/contract/test_policy_v2_compatibility.py @@ -15,7 +15,6 @@ ) from policyengine_api.services.v2.user_policies.types import LegacyUserPolicySnapshot - POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") USER_ID = UUID("00000000-0000-0000-0000-000000000070") diff --git a/tests/contract/test_simulation_gateway_contract.py b/tests/contract/test_simulation_gateway_contract.py index 2452e3e38..c890d4a88 100644 --- a/tests/contract/test_simulation_gateway_contract.py +++ b/tests/contract/test_simulation_gateway_contract.py @@ -1,10 +1,10 @@ from unittest.mock import MagicMock import httpx -import policyengine_api.libs.simulation_entrypoint as simulation_entrypoint import pytest -from policyengine_api.libs.simulation_entrypoint import SimulationAPIModal +import policyengine_api.libs.simulation_entrypoint as simulation_entrypoint +from policyengine_api.libs.simulation_entrypoint import SimulationAPIModal from tests.fixtures.libs.simulation_entrypoint import ( MOCK_BATCH_JOB_ID, MOCK_BATCH_POLL_RESPONSE_COMPLETE, diff --git a/tests/contract/test_v1_route_contracts.py b/tests/contract/test_v1_route_contracts.py index b4dce5501..5fa890323 100644 --- a/tests/contract/test_v1_route_contracts.py +++ b/tests/contract/test_v1_route_contracts.py @@ -1,6 +1,6 @@ -from contextlib import ExitStack import importlib import sys +from contextlib import ExitStack from types import SimpleNamespace from unittest.mock import patch @@ -16,25 +16,25 @@ Simulation, UserPolicy, ) -from policyengine_api.services.v2.user_policies.types import LegacyUserPolicySnapshot from policyengine_api.extensions import cache from policyengine_api.routes.household_routes import household_bp from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.routes.report_output_routes import report_output_bp from policyengine_api.routes.simulation_routes import simulation_bp -from policyengine_api.services.report_output_service import ( - ReportCreateResult, - ReportOutputView, -) from policyengine_api.services.household_calculation_service import ( HouseholdCalculationResult, ) from policyengine_api.services.household_service import HouseholdCreateResult +from policyengine_api.services.report_output_service import ( + ReportCreateResult, + ReportOutputView, +) from policyengine_api.services.simulation_service import SimulationCreateResult from policyengine_api.services.user_policy_service import ( UserPolicyCreateResult, UserPolicyUpdateResult, ) +from policyengine_api.services.v2.user_policies.types import LegacyUserPolicySnapshot from tests.contract.clients import ( ASGIContractClient, ContractClient, @@ -47,7 +47,6 @@ ) from tests.contract.registry import APP_V1_COMPATIBLE_ROUTE_CONTRACTS, ContractRequest - _STAGE12_PRIVATE_FIELDS = frozenset( { "evaluation_id", diff --git a/tests/env_variables/test_environment_variables.py b/tests/env_variables/test_environment_variables.py index 5dde7a2d2..631f6baca 100644 --- a/tests/env_variables/test_environment_variables.py +++ b/tests/env_variables/test_environment_variables.py @@ -1,4 +1,5 @@ import os + import pytest import requests diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index b0fe44636..22845e338 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -70,5 +70,4 @@ def create_household_with_axes(base_household, axes_config): def setup_small_axes_household(base_household, small_axes_config): """Fixture to setup a household with small axes for testing""" household_with_axes = create_household_with_axes(base_household, small_axes_config) - household_with_axes = add_yearly_variables(household_with_axes, TEST_COUNTRY_ID) - return household_with_axes + return add_yearly_variables(household_with_axes, TEST_COUNTRY_ID) diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index 02a444be8..de5e508f2 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -5,19 +5,20 @@ the Modal simulation API client. """ -import pytest -from unittest.mock import patch, MagicMock import json +from unittest.mock import MagicMock, patch + +import pytest from policyengine_api.constants import ( - MODAL_EXECUTION_STATUS_SUBMITTED, - MODAL_EXECUTION_STATUS_RUNNING, MODAL_EXECUTION_STATUS_COMPLETE, MODAL_EXECUTION_STATUS_FAILED, + MODAL_EXECUTION_STATUS_RUNNING, + MODAL_EXECUTION_STATUS_SUBMITTED, ) from tests.fixtures.spm import ( - worker_spm_capability, worker_result_fields, + worker_spm_capability, worker_versions_document, ) diff --git a/tests/fixtures/local_v1_database.py b/tests/fixtures/local_v1_database.py index 826d730cf..d4712f94f 100644 --- a/tests/fixtures/local_v1_database.py +++ b/tests/fixtures/local_v1_database.py @@ -1,10 +1,9 @@ """SQLite-only v1 schema helper confined to explicit test fixtures.""" -from sqlalchemy import Column, Engine, Integer, JSON, MetaData, String, Table +from sqlalchemy import JSON, Column, Engine, Integer, MetaData, String, Table from policyengine_api.data.v1_models import Policy, V1Base - _sqlite_policy_metadata = MetaData() Table( Policy.__tablename__, diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index e1ae2d190..4f3738d96 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -4,6 +4,7 @@ from uuid import UUID import pytest + from policyengine_api.constants import ( MODAL_EXECUTION_STATUS_RUNNING, MODAL_EXECUTION_STATUS_SUBMITTED, diff --git a/tests/fixtures/services/household_fixtures.py b/tests/fixtures/services/household_fixtures.py index 23602dbdf..9daa799cc 100644 --- a/tests/fixtures/services/household_fixtures.py +++ b/tests/fixtures/services/household_fixtures.py @@ -1,7 +1,8 @@ -import pytest import json from unittest.mock import patch +import pytest + from policyengine_api.data.v1_models import Household valid_request_body = { diff --git a/tests/fixtures/services/policy_service.py b/tests/fixtures/services/policy_service.py index 9b655232c..23dc38648 100644 --- a/tests/fixtures/services/policy_service.py +++ b/tests/fixtures/services/policy_service.py @@ -1,7 +1,8 @@ -import pytest import json from unittest.mock import patch +import pytest + from policyengine_api.data.v1_models import Policy valid_policy_json = { diff --git a/tests/fixtures/spm.py b/tests/fixtures/spm.py index c295633ec..118964fd0 100644 --- a/tests/fixtures/spm.py +++ b/tests/fixtures/spm.py @@ -5,7 +5,6 @@ from policyengine_api import spm from policyengine_api.constants import POLICYENGINE_VERSION - SPM_CONTRACT_VERSION = "canonical-spm-v1" INSTALLED_SPM_SELECTION = spm.normalize_spm_selection("us", None) diff --git a/tests/fixtures/v2_catalog.py b/tests/fixtures/v2_catalog.py index ad887af56..23ecff6cd 100644 --- a/tests/fixtures/v2_catalog.py +++ b/tests/fixtures/v2_catalog.py @@ -2,13 +2,12 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from types import SimpleNamespace from policyengine_api.data.v2.catalog.extraction import extract_catalog from policyengine_api.data.v2.catalog.records import NormalizedCatalog - POLICYENGINE_VERSION = "5.0.4" DEPENDENCY_VERSIONS = { "policyengine-core": "3.30.1", @@ -116,12 +115,12 @@ def _model_source( parameter_values=[ SimpleNamespace( value=0.1, - start_date=datetime(2025, 1, 1, tzinfo=timezone.utc), - end_date=datetime(2025, 12, 31, tzinfo=timezone.utc), + start_date=datetime(2025, 1, 1, tzinfo=UTC), + end_date=datetime(2025, 12, 31, tzinfo=UTC), ), SimpleNamespace( value=0.2, - start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), + start_date=datetime(2026, 1, 1, tzinfo=UTC), end_date=None, ), ], diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index b9f61c8c1..a4935ec1c 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -2,20 +2,20 @@ import os +import pytest from alembic import command from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext -import pytest from sqlalchemy import ( Column, - delete, Integer, MetaData, String, Table, Text, create_engine, + delete, inspect, ) from sqlalchemy.dialects.mysql import LONGTEXT @@ -37,7 +37,6 @@ upgrade_database, ) - BASELINE_REVISION = "eafc2a547a4e" PREVIOUS_REVISION = "3d6e8f553ca5" diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 7c6628a48..66120c3da 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -3,12 +3,12 @@ import os from uuid import UUID, uuid4 +import pytest +import sqlalchemy as sa from alembic import command from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext -import pytest -import sqlalchemy as sa from sqlalchemy import create_engine, inspect, text from sqlalchemy.dialects.postgresql import UUID as PostgresUUID @@ -22,7 +22,6 @@ from policyengine_api.data.v2.models import V2_METADATA from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL - BASELINE_REVISION = "f5ef4347cb2a" STAGE_11_PREVIOUS_REVISION = "af34023a728f" STAGE_12_PREVIOUS_REVISION = "724b1b11a33e" diff --git a/tests/integration/test_budget_window_in_flight_dedupe.py b/tests/integration/test_budget_window_in_flight_dedupe.py index 826154ca9..b6c5c64bd 100644 --- a/tests/integration/test_budget_window_in_flight_dedupe.py +++ b/tests/integration/test_budget_window_in_flight_dedupe.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock from flask import Flask + from policyengine_api.runtime_cache.fake import InMemoryCacheBackend from tests.fixtures.spm import worker_spm_capability diff --git a/tests/integration/test_live_economy.py b/tests/integration/test_live_economy.py index 259f7949a..5e9d062c4 100644 --- a/tests/integration/test_live_economy.py +++ b/tests/integration/test_live_economy.py @@ -1,6 +1,6 @@ -from datetime import date import json import math +from datetime import date from pathlib import Path CURRENT_YEAR = date.today().year diff --git a/tests/integration/test_live_v2_metadata.py b/tests/integration/test_live_v2_metadata.py index 466745e17..235304156 100644 --- a/tests/integration/test_live_v2_metadata.py +++ b/tests/integration/test_live_v2_metadata.py @@ -1,11 +1,10 @@ """Read-only checks for deployed Cloud Run v2 metadata preview routes.""" -from pathlib import Path import tomllib +from pathlib import Path import pytest - REPO = Path(__file__).parents[2] PAGED_RESOURCES = ( "tax-benefit-models", diff --git a/tests/integration/test_mysql_household_spm_dual_write.py b/tests/integration/test_mysql_household_spm_dual_write.py index ec42e3485..20e749715 100644 --- a/tests/integration/test_mysql_household_spm_dual_write.py +++ b/tests/integration/test_mysql_household_spm_dual_write.py @@ -7,10 +7,10 @@ from __future__ import annotations +import json from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass -import json from uuid import uuid4 import pytest @@ -22,6 +22,8 @@ from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS from policyengine_api.data.v1_models import ( Household as V1Household, +) +from policyengine_api.data.v1_models import ( HouseholdMirrorEvent, ) from policyengine_api.data.v2.models import Household, LegacyHouseholdMapping @@ -41,7 +43,6 @@ from policyengine_api.utils import hash_object from tests.integration.test_mysql_policy_dual_write import _mysql_url, _postgres_url - NATIONAL_SPM = { "forecast_content_sha256": "a" * 64, "scenario": "baseline", diff --git a/tests/integration/test_runtime_cache_redis.py b/tests/integration/test_runtime_cache_redis.py index e12b04575..ba332c7b7 100644 --- a/tests/integration/test_runtime_cache_redis.py +++ b/tests/integration/test_runtime_cache_redis.py @@ -1,8 +1,8 @@ """Real Redis qualification for cross-connection, TTL, atomicity, and claims.""" -from datetime import datetime import os import time +from datetime import datetime from urllib.parse import urlsplit from uuid import uuid4 @@ -22,7 +22,6 @@ reform_impact_id, ) - RUNTIME_CACHE_TEST_URL = "RUNTIME_CACHE_TEST_URL" diff --git a/tests/integration/test_simulations.py b/tests/integration/test_simulations.py index bcb3c5d64..7bc19e2e3 100644 --- a/tests/integration/test_simulations.py +++ b/tests/integration/test_simulations.py @@ -1,12 +1,12 @@ from policyengine_api.country import COUNTRIES from tests.fixtures.integration.simulations import ( + SMALL_AXES_COUNT, TEST_COUNTRY_ID, TEST_STATE, TEST_YEAR, - SMALL_AXES_COUNT, - setup_small_axes_household, create_base_household, create_small_axes, + setup_small_axes_household, ) diff --git a/tests/integration/test_v1_household_dual_write.py b/tests/integration/test_v1_household_dual_write.py index dba2573da..f36de6a9f 100644 --- a/tests/integration/test_v1_household_dual_write.py +++ b/tests/integration/test_v1_household_dual_write.py @@ -12,6 +12,8 @@ from policyengine_api.data.v1_models import ( Household as V1Household, +) +from policyengine_api.data.v1_models import ( HouseholdMirrorEvent, ) from policyengine_api.data.v2.models import Household, LegacyHouseholdMapping diff --git a/tests/integration/test_v1_policy_dual_write.py b/tests/integration/test_v1_policy_dual_write.py index d23d5ecef..4148fda52 100644 --- a/tests/integration/test_v1_policy_dual_write.py +++ b/tests/integration/test_v1_policy_dual_write.py @@ -18,6 +18,11 @@ TaxBenefitModel, TaxBenefitModelVersion, ) +from policyengine_api.services.policy_mirroring import ( + PolicyMirrorUnavailableError, + mirror_policy_after_commit, +) +from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.v2.policies.database_session import ( PolicyDatabaseSession, ) @@ -25,11 +30,6 @@ V2PolicyService, mirror_legacy_policy_in_session, ) -from policyengine_api.services.policy_mirroring import ( - PolicyMirrorUnavailableError, - mirror_policy_after_commit, -) -from policyengine_api.services.policy_service import PolicyService def _v1_service(database_url: str): diff --git a/tests/integration/test_v1_user_policy_dual_write.py b/tests/integration/test_v1_user_policy_dual_write.py index a9679bc54..2bcd63e2a 100644 --- a/tests/integration/test_v1_user_policy_dual_write.py +++ b/tests/integration/test_v1_user_policy_dual_write.py @@ -10,7 +10,11 @@ from policyengine_api.constants import POLICYENGINE_VERSION from policyengine_api.data.v1_models import ( Policy as V1Policy, +) +from policyengine_api.data.v1_models import ( UserPolicy as V1UserPolicy, +) +from policyengine_api.data.v1_models import ( UserPolicyMirrorEvent, ) from policyengine_api.data.v2.models import ( @@ -25,10 +29,6 @@ User, UserPolicy, ) -from policyengine_api.services.v2.user_policies.database_session import ( - UserPolicyDatabaseSession, -) -from policyengine_api.services.v2.user_policies.services import V2UserPolicyService from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.user_policy_mirroring import ( UserPolicyMirrorUnavailableError, @@ -36,6 +36,10 @@ mirror_user_policy_after_commit, ) from policyengine_api.services.user_policy_service import UserPolicyService +from policyengine_api.services.v2.user_policies.database_session import ( + UserPolicyDatabaseSession, +) +from policyengine_api.services.v2.user_policies.services import V2UserPolicyService def _v1_services(database_url: str): diff --git a/tests/integration/test_v2_catalog_installed.py b/tests/integration/test_v2_catalog_installed.py index 06f7e79e2..c2f7dfae8 100644 --- a/tests/integration/test_v2_catalog_installed.py +++ b/tests/integration/test_v2_catalog_installed.py @@ -2,15 +2,14 @@ from __future__ import annotations -from importlib import metadata as importlib_metadata import json import os +from importlib import metadata as importlib_metadata import pytest - from policyengine.bundle import get_current_bundle -from policyengine_api.data.v2.catalog.extraction import extract_installed_catalog +from policyengine_api.data.v2.catalog.extraction import extract_installed_catalog pytestmark = pytest.mark.skipif( os.environ.get("RUN_V2_CATALOG_COMPATIBILITY") != "1", diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index e5d21f399..cb59127ce 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -2,14 +2,14 @@ from __future__ import annotations +import os from concurrent.futures import ThreadPoolExecutor from dataclasses import replace -import os from pathlib import Path +import pytest from alembic import command from alembic.config import Config -import pytest from sqlalchemy import create_engine, text from sqlalchemy.engine import Engine, make_url from sqlalchemy.pool import NullPool @@ -19,7 +19,6 @@ CatalogPublicationError, publish_catalog, ) -from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL from policyengine_api.data.v2.models import ( Dataset, DatasetVersion, @@ -30,9 +29,9 @@ TaxBenefitModel, TaxBenefitModelVersion, ) +from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL from tests.fixtures.v2_catalog import normalized_catalog - REPO = Path(__file__).parents[2] DISPOSABLE_DATABASE = "policyengine_v2_alembic_test" diff --git a/tests/integration/test_v2_catalog_publication_qualification.py b/tests/integration/test_v2_catalog_publication_qualification.py index 2a0fe1bde..8400c4ac9 100644 --- a/tests/integration/test_v2_catalog_publication_qualification.py +++ b/tests/integration/test_v2_catalog_publication_qualification.py @@ -5,12 +5,12 @@ import gc import json import os -from pathlib import Path import tracemalloc +from pathlib import Path +import pytest from alembic import command from alembic.config import Config -import pytest from sqlalchemy import create_engine, text from sqlalchemy.engine import make_url from sqlalchemy.pool import NullPool @@ -19,7 +19,6 @@ from policyengine_api.data.v2.catalog.publication import publish_catalog from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL - REPO = Path(__file__).parents[2] DISPOSABLE_DATABASE = "policyengine_v2_alembic_test" MAX_ADDITIONAL_PUBLISHER_MEMORY_BYTES = 256 * 1024 * 1024 diff --git a/tests/integration/test_v2_household_persistence.py b/tests/integration/test_v2_household_persistence.py index da1ccbe08..a4f50801c 100644 --- a/tests/integration/test_v2_household_persistence.py +++ b/tests/integration/test_v2_household_persistence.py @@ -2,8 +2,8 @@ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor import os +from concurrent.futures import ThreadPoolExecutor from threading import Barrier from uuid import UUID, uuid4 diff --git a/tests/integration/test_v2_metadata_routes.py b/tests/integration/test_v2_metadata_routes.py index f1031637a..a3d29d93b 100644 --- a/tests/integration/test_v2_metadata_routes.py +++ b/tests/integration/test_v2_metadata_routes.py @@ -6,11 +6,11 @@ from pathlib import Path from uuid import uuid4 +import pytest from alembic import command from alembic.config import Config from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from sqlalchemy import create_engine, text from sqlalchemy.engine import Engine, make_url from sqlalchemy.pool import NullPool @@ -18,10 +18,6 @@ from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.data.v2.catalog.publication import publish_catalog -from policyengine_api.services.v2.metadata.database_session import ( - MetadataDatabaseSession, -) -from policyengine_api.services.v2.metadata.services import V2MetadataService from policyengine_api.data.v2.models import ( Dataset, TaxBenefitModel, @@ -33,9 +29,12 @@ RouteImplementation, RouteImplementationSettings, ) +from policyengine_api.services.v2.metadata.database_session import ( + MetadataDatabaseSession, +) +from policyengine_api.services.v2.metadata.services import V2MetadataService from tests.fixtures.v2_catalog import POLICYENGINE_VERSION, normalized_catalog - REPO = Path(__file__).parents[2] DISPOSABLE_DATABASE = "policyengine_v2_alembic_test" @@ -80,7 +79,7 @@ def v1_metadata(country_id: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_metadata_reader_factory=lambda: V2MetadataService( MetadataDatabaseSession(Session(engine)), running_policyengine_version=POLICYENGINE_VERSION, diff --git a/tests/integration/test_v2_policy_persistence.py b/tests/integration/test_v2_policy_persistence.py index e5cad92b2..4ca019529 100644 --- a/tests/integration/test_v2_policy_persistence.py +++ b/tests/integration/test_v2_policy_persistence.py @@ -2,8 +2,8 @@ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor import os +from concurrent.futures import ThreadPoolExecutor from threading import Barrier from uuid import UUID, uuid4 @@ -24,14 +24,15 @@ TaxBenefitModel, TaxBenefitModelVersion, ) -from policyengine_api.services.v2.policies.transformations import ( - canonical_policy_document, - canonicalize_policy, -) +from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL from policyengine_api.services.v2.policies.services import ( create_resolved_policy, mirror_legacy_policy_in_session, ) +from policyengine_api.services.v2.policies.transformations import ( + canonical_policy_document, + canonicalize_policy, +) from policyengine_api.services.v2.policies.types import ( CanonicalPolicyContent, LegacyPolicySnapshot, @@ -41,7 +42,6 @@ LegacyPolicyMappingIntegrityError, PolicyContentHashCollisionError, ) -from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL def _disposable_url() -> str: diff --git a/tests/integration/test_v2_user_policy_mirroring.py b/tests/integration/test_v2_user_policy_mirroring.py index 1b691451f..c2819f716 100644 --- a/tests/integration/test_v2_user_policy_mirroring.py +++ b/tests/integration/test_v2_user_policy_mirroring.py @@ -2,8 +2,8 @@ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor import os +from concurrent.futures import ThreadPoolExecutor from threading import Barrier from uuid import UUID, uuid4 diff --git a/tests/to_refactor/api/test_api.py b/tests/to_refactor/api/test_api.py index 68b4d5b04..7fcb7f705 100644 --- a/tests/to_refactor/api/test_api.py +++ b/tests/to_refactor/api/test_api.py @@ -1,7 +1,9 @@ -import yaml import json -import pytest from pathlib import Path + +import pytest +import yaml + from policyengine_api.api import app diff --git a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py index 8dbd4bb23..840ebe312 100644 --- a/tests/to_refactor/fixtures/to_refactor_household_fixtures.py +++ b/tests/to_refactor/fixtures/to_refactor_household_fixtures.py @@ -1,7 +1,8 @@ import json -import pytest from unittest.mock import patch +import pytest + valid_request_body = { "data": {"people": {"person1": {"age": 30, "income": 50000}}}, "label": "Test Household", diff --git a/tests/to_refactor/python/test_error_routes.py b/tests/to_refactor/python/test_error_routes.py index c36d75b24..09c2058da 100644 --- a/tests/to_refactor/python/test_error_routes.py +++ b/tests/to_refactor/python/test_error_routes.py @@ -1,14 +1,15 @@ import pytest from flask import Flask -from policyengine_api.routes.error_routes import error_bp from werkzeug.exceptions import ( - NotFound, BadRequest, - Unauthorized, Forbidden, InternalServerError, + NotFound, + Unauthorized, ) +from policyengine_api.routes.error_routes import error_bp + @pytest.fixture def app(): diff --git a/tests/to_refactor/python/test_household_routes.py b/tests/to_refactor/python/test_household_routes.py index fd57f09eb..858c20cb3 100644 --- a/tests/to_refactor/python/test_household_routes.py +++ b/tests/to_refactor/python/test_household_routes.py @@ -4,8 +4,8 @@ from policyengine_api.data.v1_models import Household from policyengine_api.services.household_service import HouseholdCreateResult from tests.to_refactor.fixtures.to_refactor_household_fixtures import ( - valid_request_body, valid_db_row, + valid_request_body, ) pytest_plugins = ["tests.to_refactor.fixtures.to_refactor_household_fixtures"] diff --git a/tests/to_refactor/python/test_policy.py b/tests/to_refactor/python/test_policy.py index 515634dc0..b6adf8397 100644 --- a/tests/to_refactor/python/test_policy.py +++ b/tests/to_refactor/python/test_policy.py @@ -1,4 +1,5 @@ import json + from sqlalchemy import delete from policyengine_api.data.orm import get_v1_session_factory @@ -17,7 +18,7 @@ class TestPolicyCreation: """ Test creating a policy, then ensure that duplicating that policy generates the correct response within the - app; this requires sequential processing, hence the + app; this requires sequential processing, hence the need for a separate Python-based test """ diff --git a/tests/to_refactor/python/test_validate_country.py b/tests/to_refactor/python/test_validate_country.py index d927198aa..f3495ffe8 100644 --- a/tests/to_refactor/python/test_validate_country.py +++ b/tests/to_refactor/python/test_validate_country.py @@ -1,4 +1,5 @@ from flask import Response + from policyengine_api.utils.payload_validators import validate_country diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index b72075604..7c8bd744e 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,7 +4,6 @@ from sqlalchemy import create_engine from sqlalchemy.pool import StaticPool - os.environ.setdefault("FLASK_DEBUG", "1") from policyengine_api.data import orm diff --git a/tests/unit/data/test_alembic_baseline.py b/tests/unit/data/test_alembic_baseline.py index a6bf8aeb5..937b3c2f7 100644 --- a/tests/unit/data/test_alembic_baseline.py +++ b/tests/unit/data/test_alembic_baseline.py @@ -1,9 +1,9 @@ from io import StringIO +import pytest from alembic import command from alembic.config import Config from alembic.script import ScriptDirectory -import pytest from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base diff --git a/tests/unit/data/test_congressional_districts.py b/tests/unit/data/test_congressional_districts.py index 7fd1bfb42..e2a819793 100644 --- a/tests/unit/data/test_congressional_districts.py +++ b/tests/unit/data/test_congressional_districts.py @@ -2,13 +2,13 @@ from pydantic import ValidationError from policyengine_api.data.congressional_districts import ( - CongressionalDistrictMetadataItem, + AT_LARGE_STATES, CONGRESSIONAL_DISTRICTS, STATE_CODE_TO_NAME, - AT_LARGE_STATES, + CongressionalDistrictMetadataItem, build_congressional_district_metadata, - get_valid_state_codes, get_valid_congressional_districts, + get_valid_state_codes, normalize_us_region, ) diff --git a/tests/unit/data/test_runtime_sql_boundaries.py b/tests/unit/data/test_runtime_sql_boundaries.py index 7b447aa9d..4cf05b526 100644 --- a/tests/unit/data/test_runtime_sql_boundaries.py +++ b/tests/unit/data/test_runtime_sql_boundaries.py @@ -2,7 +2,6 @@ from pathlib import Path - PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 59dad5287..095ace664 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/tests/unit/data/test_sqlalchemy_v2.py @@ -1,7 +1,7 @@ +from pathlib import Path from unittest.mock import Mock import sqlalchemy -from pathlib import Path import policyengine_api.data.orm as orm diff --git a/tests/unit/data/test_v1_models.py b/tests/unit/data/test_v1_models.py index bb3111e59..da5476fe7 100644 --- a/tests/unit/data/test_v1_models.py +++ b/tests/unit/data/test_v1_models.py @@ -3,7 +3,6 @@ from policyengine_api.data.v1_models import V1Base - EXPECTED_TABLES = { "analysis", "computed_household", diff --git a/tests/unit/libs/test_gateway_auth.py b/tests/unit/libs/test_gateway_auth.py index f9e1d8ac8..8103c9a69 100644 --- a/tests/unit/libs/test_gateway_auth.py +++ b/tests/unit/libs/test_gateway_auth.py @@ -22,7 +22,6 @@ gateway_auth_required, ) - ISSUER = "https://policyengine.uk.auth0.com" AUDIENCE = "https://sim-gateway.policyengine.org" CLIENT_ID = "test-client-id" diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 833c85f68..5616b8dae 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -15,11 +15,11 @@ os.environ.setdefault("FLASK_DEBUG", "1") from policyengine_api.constants import ( # noqa: E402 - POLICYENGINE_VERSION, MODAL_EXECUTION_STATUS_COMPLETE, MODAL_EXECUTION_STATUS_FAILED, MODAL_EXECUTION_STATUS_RUNNING, MODAL_EXECUTION_STATUS_SUBMITTED, + POLICYENGINE_VERSION, ) from policyengine_api.libs.simulation_entrypoint import ( # noqa: E402 ModalBudgetWindowBatchExecution, @@ -28,6 +28,9 @@ SimulationEntrypointClient, resolve_simulation_entrypoint_url, ) +from policyengine_api.observability.identifiers import ( # noqa: E402 + OBSERVABILITY_ID_HEADER, +) from policyengine_api.request_context import ( # noqa: E402 REQUEST_ID_HEADER, _asgi_observability_id, @@ -35,10 +38,6 @@ current_observability_id, start_observability_id, ) -from policyengine_api.observability.identifiers import ( # noqa: E402 - OBSERVABILITY_ID_HEADER, -) - from tests.fixtures.libs.simulation_entrypoint import ( # noqa: E402 MOCK_BATCH_JOB_ID, MOCK_BATCH_POLL_RESPONSE_COMPLETE, @@ -48,12 +47,12 @@ MOCK_HEALTH_RESPONSE, MOCK_MODAL_BASE_URL, MOCK_MODAL_JOB_ID, + MOCK_OBSERVABILITY_ID, MOCK_POLICYENGINE_BUNDLE, MOCK_POLL_RESPONSE_COMPLETE, MOCK_POLL_RESPONSE_FAILED, MOCK_POLL_RESPONSE_RUNNING, MOCK_RESOLVED_APP_NAME, - MOCK_OBSERVABILITY_ID, MOCK_SIMULATION_PAYLOAD, MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY, MOCK_SIMULATION_RESULT, diff --git a/tests/unit/routes/test_calculate_deprecated_inputs.py b/tests/unit/routes/test_calculate_deprecated_inputs.py index 3c52df5b4..4081d30cc 100644 --- a/tests/unit/routes/test_calculate_deprecated_inputs.py +++ b/tests/unit/routes/test_calculate_deprecated_inputs.py @@ -1,9 +1,10 @@ -from flask import Flask -import pytest from unittest.mock import patch -from policyengine_api.routes import household_routes +import pytest +from flask import Flask + from policyengine_api.extensions import cache +from policyengine_api.routes import household_routes from policyengine_api.services.household_calculation_service import ( CalculationResult, HouseholdCalculationService, diff --git a/tests/unit/routes/test_calculate_error_statuses.py b/tests/unit/routes/test_calculate_error_statuses.py index 4af690247..5972db9d3 100644 --- a/tests/unit/routes/test_calculate_error_statuses.py +++ b/tests/unit/routes/test_calculate_error_statuses.py @@ -1,6 +1,7 @@ +from unittest.mock import patch + from flask import Flask from policyengine_core.errors import SituationParsingError -from unittest.mock import patch from policyengine_api.extensions import cache from policyengine_api.routes import household_routes diff --git a/tests/unit/routes/test_canonical_spm.py b/tests/unit/routes/test_canonical_spm.py index 7e8c678e8..3c7095625 100644 --- a/tests/unit/routes/test_canonical_spm.py +++ b/tests/unit/routes/test_canonical_spm.py @@ -4,8 +4,8 @@ from types import SimpleNamespace from unittest.mock import patch -from flask import Flask import pytest +from flask import Flask from policyengine_api import spm from policyengine_api.data.v1_models import Household, Policy, Simulation @@ -24,7 +24,6 @@ from policyengine_api.services.simulation_service import SimulationService from policyengine_api.utils import hash_object - HOUSEHOLD = {"people": {"you": {"age": {"2026": 40}}}} FORECAST_HASH = "a" * 64 @@ -320,7 +319,7 @@ def test_certification_checked_before_cached_response(certified, harness): @pytest.mark.parametrize("country_id", ["us", "uk"]) def test_legacy_country_requests_do_not_receive_spm(harness, country_id, monkeypatch): - monkeypatch.setattr(spm, "_current_bundle", lambda: {}) + monkeypatch.setattr(spm, "_current_bundle", dict) monkeypatch.setattr(spm, "simulation_supports_spm", lambda _: False) client, country = harness response = client.post(f"/{country_id}/calculate", json={"household": HOUSEHOLD}) @@ -446,10 +445,12 @@ def test_certification_protects_linked_households_without_saved_spm( def test_household_creation_retains_resolved_spm_before_selected_mirror( certified, harness, orm_session_factory, monkeypatch, write_source ): - from policyengine_api.data.v1_models import HouseholdMirrorEvent - from sqlalchemy import select from unittest.mock import Mock + from sqlalchemy import select + + from policyengine_api.data.v1_models import HouseholdMirrorEvent + monkeypatch.setenv("DB_WRITE_HOUSEHOLD", write_source) observed = [] @@ -521,9 +522,10 @@ def test_invalid_spm_stops_before_database_and_mirror( def test_valid_spm_database_timeout_keeps_stage11_safe_persistence_response( certified, harness, monkeypatch, write_source ): - from sqlalchemy.exc import TimeoutError from unittest.mock import Mock + from sqlalchemy.exc import TimeoutError + sessions = Mock() sessions.begin.side_effect = TimeoutError("database credential secret") monkeypatch.setenv("DB_WRITE_HOUSEHOLD", write_source) diff --git a/tests/unit/routes/test_economy_failure_replay.py b/tests/unit/routes/test_economy_failure_replay.py index a1889f603..13ce48997 100644 --- a/tests/unit/routes/test_economy_failure_replay.py +++ b/tests/unit/routes/test_economy_failure_replay.py @@ -3,9 +3,10 @@ from types import SimpleNamespace from unittest.mock import Mock -from flask import Flask import pytest +from flask import Flask +from policyengine_api import worker_spm from policyengine_api.routes import economy_routes from policyengine_api.runtime_cache.core import CacheNamespace from policyengine_api.runtime_cache.fake import InMemoryCacheBackend @@ -13,8 +14,6 @@ from policyengine_api.services import economy_service as economy_module from policyengine_api.services.economy_service import EconomyService from policyengine_api.services.reform_impacts_service import ReformImpactsService -from policyengine_api import worker_spm - SELECTION = { "forecast_content_sha256": "a" * 64, diff --git a/tests/unit/routes/test_economy_spm_query.py b/tests/unit/routes/test_economy_spm_query.py index 7c6d2e598..90a14ced2 100644 --- a/tests/unit/routes/test_economy_spm_query.py +++ b/tests/unit/routes/test_economy_spm_query.py @@ -4,8 +4,8 @@ from types import SimpleNamespace from unittest.mock import Mock -from flask import Flask import pytest +from flask import Flask from policyengine_api import spm from policyengine_api.routes import economy_routes diff --git a/tests/unit/routes/test_flask_route_boundaries.py b/tests/unit/routes/test_flask_route_boundaries.py index ae2f01b60..2251d46d2 100644 --- a/tests/unit/routes/test_flask_route_boundaries.py +++ b/tests/unit/routes/test_flask_route_boundaries.py @@ -8,7 +8,6 @@ from policyengine_api.routes.reform_impact_routes import reform_impact_bp from policyengine_api.routes.system_routes import system_bp - PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" diff --git a/tests/unit/routes/test_household_and_user_policy_orm_routes.py b/tests/unit/routes/test_household_and_user_policy_orm_routes.py index 43c782c52..46efdcd1b 100644 --- a/tests/unit/routes/test_household_and_user_policy_orm_routes.py +++ b/tests/unit/routes/test_household_and_user_policy_orm_routes.py @@ -11,6 +11,12 @@ Policy, UserPolicy, ) +from policyengine_api.routes.household_routes import get_household_under_policy +from policyengine_api.routes.policy_routes import ( + get_user_policy, + set_user_policy, + update_user_policy, +) from policyengine_api.runtime_cache.core import CacheNamespace from policyengine_api.runtime_cache.fake import InMemoryCacheBackend from policyengine_api.runtime_cache.household_calculations import ( @@ -18,12 +24,6 @@ HouseholdCalculationCache, HouseholdCalculationIdentity, ) -from policyengine_api.routes.household_routes import get_household_under_policy -from policyengine_api.routes.policy_routes import ( - get_user_policy, - set_user_policy, - update_user_policy, -) from policyengine_api.services.household_calculation_service import ( CalculationResult, HouseholdCalculationService, diff --git a/tests/unit/routes/test_household_dual_write_routes.py b/tests/unit/routes/test_household_dual_write_routes.py index 84c1b1ba6..bc94a8487 100644 --- a/tests/unit/routes/test_household_dual_write_routes.py +++ b/tests/unit/routes/test_household_dual_write_routes.py @@ -4,8 +4,8 @@ from unittest.mock import MagicMock, patch -from flask import Flask import pytest +from flask import Flask from policyengine_api.data.v1_models import Household from policyengine_api.routes.household_routes import household_bp diff --git a/tests/unit/routes/test_household_numeric_identity.py b/tests/unit/routes/test_household_numeric_identity.py index e65e4b13b..604c5f710 100644 --- a/tests/unit/routes/test_household_numeric_identity.py +++ b/tests/unit/routes/test_household_numeric_identity.py @@ -1,10 +1,10 @@ """Immutable households retain strict numeric simulation alias boundaries.""" -from copy import deepcopy import json +from copy import deepcopy -from flask import Flask import pytest +from flask import Flask from policyengine_api.data.v1_models import Simulation, SimulationRun from policyengine_api.routes import household_routes, simulation_routes @@ -12,7 +12,6 @@ from policyengine_api.services.household_service import HouseholdService from policyengine_api.services.simulation_service import SimulationService - INPUTS = {"people": {"you": {"age": {"2026": 40}}}} NATIONAL = { "forecast_content_sha256": "a" * 64, diff --git a/tests/unit/routes/test_policy_dual_write_routes.py b/tests/unit/routes/test_policy_dual_write_routes.py index 39b2e1771..4dfd49cca 100644 --- a/tests/unit/routes/test_policy_dual_write_routes.py +++ b/tests/unit/routes/test_policy_dual_write_routes.py @@ -5,14 +5,14 @@ import json from unittest.mock import MagicMock, patch -from flask import Flask import pytest +from flask import Flask from policyengine_api.data.v1_models import Policy -from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.services.policy_mirroring import PolicyMirrorUnavailableError from policyengine_api.services.policy_service import PolicySetResult +from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot def _client(): diff --git a/tests/unit/routes/test_reform_impact_routes.py b/tests/unit/routes/test_reform_impact_routes.py index 422db041a..29ffbf65e 100644 --- a/tests/unit/routes/test_reform_impact_routes.py +++ b/tests/unit/routes/test_reform_impact_routes.py @@ -5,12 +5,12 @@ v1 response-shape contracts remain unchanged. """ -from datetime import datetime import json +from datetime import datetime +import pytest from fastapi.testclient import TestClient from flask import Flask -import pytest from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.routes import reform_impact_routes @@ -23,7 +23,6 @@ ) from policyengine_api.services.reform_impacts_service import ReformImpactsService - COMPLETED_AT = datetime(2026, 1, 1, 1) diff --git a/tests/unit/routes/test_report_household_identity.py b/tests/unit/routes/test_report_household_identity.py index 36ae50b8c..4361e3bbd 100644 --- a/tests/unit/routes/test_report_household_identity.py +++ b/tests/unit/routes/test_report_household_identity.py @@ -2,8 +2,8 @@ import json -from flask import Flask import pytest +from flask import Flask from sqlalchemy import select from policyengine_api.data.v1_models import ( diff --git a/tests/unit/routes/test_spm_year_worker_polling.py b/tests/unit/routes/test_spm_year_worker_polling.py index 3faa0ab02..ca323ea67 100644 --- a/tests/unit/routes/test_spm_year_worker_polling.py +++ b/tests/unit/routes/test_spm_year_worker_polling.py @@ -1,8 +1,8 @@ """Typed worker failures become durable terminal HTTP polling responses.""" -from flask import Flask import httpx import pytest +from flask import Flask from policyengine_api import worker_spm from policyengine_api.libs.simulation_entrypoint import SimulationEntrypointClient diff --git a/tests/unit/routes/test_user_policy_dual_write_routes.py b/tests/unit/routes/test_user_policy_dual_write_routes.py index 5b9d97734..87c323325 100644 --- a/tests/unit/routes/test_user_policy_dual_write_routes.py +++ b/tests/unit/routes/test_user_policy_dual_write_routes.py @@ -5,18 +5,18 @@ import json from unittest.mock import MagicMock, patch -from flask import Flask import pytest +from flask import Flask from sqlalchemy.exc import ( IntegrityError, OperationalError, SQLAlchemyError, +) +from sqlalchemy.exc import ( TimeoutError as SQLAlchemyTimeoutError, ) from policyengine_api.data.v1_models import UserPolicy -from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot -from policyengine_api.services.v2.user_policies.types import LegacyUserPolicySnapshot from policyengine_api.routes.policy_routes import policy_bp from policyengine_api.services.user_policy_mirroring import ( UserPolicyMirrorUnavailableError, @@ -26,6 +26,8 @@ UserPolicyPersistenceError, UserPolicyUpdateResult, ) +from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot +from policyengine_api.services.v2.user_policies.types import LegacyUserPolicySnapshot def _client(): diff --git a/tests/unit/runtime_cache/test_settings.py b/tests/unit/runtime_cache/test_settings.py index e112e7342..5f7c9d699 100644 --- a/tests/unit/runtime_cache/test_settings.py +++ b/tests/unit/runtime_cache/test_settings.py @@ -12,7 +12,6 @@ load_runtime_cache_settings, ) - TEST_CA_CERT = """-----BEGIN CERTIFICATE----- stage-8-test-ca -----END CERTIFICATE-----""" diff --git a/tests/unit/services/test_budget_window_cache.py b/tests/unit/services/test_budget_window_cache.py index ff344e924..2acaf0752 100644 --- a/tests/unit/services/test_budget_window_cache.py +++ b/tests/unit/services/test_budget_window_cache.py @@ -2,8 +2,11 @@ import pytest -from policyengine_api.runtime_cache.core import CacheCoordinationError -from policyengine_api.runtime_cache.core import CacheNamespace, encode_envelope +from policyengine_api.runtime_cache.core import ( + CacheCoordinationError, + CacheNamespace, + encode_envelope, +) from policyengine_api.runtime_cache.fake import InMemoryCacheBackend from policyengine_api.services.budget_window_cache import ( BUDGET_WINDOW_BATCH_TTL_SECONDS, @@ -27,7 +30,7 @@ def __init__(self, *, method): def get(self, key): if self.method == "get": raise RuntimeError("redis unavailable") - return None + return def set(self, key, value, nx=False, ex=None): if self.method == "set": diff --git a/tests/unit/services/test_economy_comparison.py b/tests/unit/services/test_economy_comparison.py index f0e451d45..3982731b3 100644 --- a/tests/unit/services/test_economy_comparison.py +++ b/tests/unit/services/test_economy_comparison.py @@ -1,19 +1,20 @@ -import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import MagicMock, patch + import numpy as np import pandas as pd +import pytest from pydantic import ValidationError from policyengine_api.services.economy_comparison import ( - UKConstituencyBreakdownByConstituency, UKConstituencyBreakdown, - UKLocalAuthorityBreakdownByLA, + UKConstituencyBreakdownByConstituency, UKLocalAuthorityBreakdown, - uk_constituency_breakdown, - uk_local_authority_breakdown, + UKLocalAuthorityBreakdownByLA, compute_income_change, intra_decile_impact, intra_wealth_decile_impact, + uk_constituency_breakdown, + uk_local_authority_breakdown, ) @@ -741,8 +742,8 @@ def _make_economy( n = len(incomes) return { "household_net_income": np.array(incomes, dtype=float), - "household_weight": np.array(weights if weights else [1.0] * n), - "household_count_people": np.array(people if people else [1.0] * n), + "household_weight": np.array(weights or [1.0] * n), + "household_count_people": np.array(people or [1.0] * n), decile_key: np.array(deciles, dtype=float), } diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index e9c30264f..d50262bd2 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -4,23 +4,24 @@ import httpx import pytest + from policyengine_api.runtime_cache.core import CacheCoordinationError from policyengine_api.runtime_cache.reform_impacts import ReformImpactStartClaim from policyengine_api.services.budget_window_cache import BudgetWindowCacheState -from policyengine_api.services.reform_impacts_service import ( - ReformImpactHandoffError, -) from policyengine_api.services.economy_service import ( BUDGET_WINDOW_MAX_END_YEAR, BUDGET_WINDOW_MAX_YEARS, + REFORM_IMPACT_START_CLAIM_ATTEMPTS, EconomicImpactResult, EconomicImpactSetupOptions, EconomyService, ImpactAction, ImpactStatus, - REFORM_IMPACT_START_CLAIM_ATTEMPTS, ) from policyengine_api.services.policy_service import PolicyService +from policyengine_api.services.reform_impacts_service import ( + ReformImpactHandoffError, +) from policyengine_api.spm import SPMValidationError from tests.fixtures.services.economy_service import ( MOCK_API_VERSION, @@ -31,16 +32,16 @@ MOCK_EXECUTION_ID, MOCK_LOOKUP_OPTIONS_HASH, MOCK_MODEL_VERSION, + MOCK_OBSERVABILITY_ID, MOCK_OPTIONS, MOCK_OPTIONS_HASH, MOCK_POLICY_ID, MOCK_POLICYENGINE_VERSION, - MOCK_SUBMISSION_CLAIM_ID, MOCK_REFORM_IMPACT_DATA, MOCK_REGION, MOCK_RESOLVED_APP_NAME, MOCK_RESOLVED_DATASET, - MOCK_OBSERVABILITY_ID, + MOCK_SUBMISSION_CLAIM_ID, MOCK_TIME_PERIOD, create_mock_budget_window_batch_execution, create_mock_reform_impact, diff --git a/tests/unit/services/test_household_calculation_service.py b/tests/unit/services/test_household_calculation_service.py index 5199da998..6f52213d0 100644 --- a/tests/unit/services/test_household_calculation_service.py +++ b/tests/unit/services/test_household_calculation_service.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, Mock import pytest + from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, POLICYENGINE_VERSION from policyengine_api.data.v1_models import ( Household, @@ -24,7 +25,6 @@ ) from tests.fixtures.spm import INSTALLED_SPM_SELECTION, household_result_fields - PACKAGE_ROOT = Path(__file__).parents[3] / "policyengine_api" diff --git a/tests/unit/services/test_household_mirroring.py b/tests/unit/services/test_household_mirroring.py index ac821dee5..f8e26afc9 100644 --- a/tests/unit/services/test_household_mirroring.py +++ b/tests/unit/services/test_household_mirroring.py @@ -21,7 +21,6 @@ LegacyHouseholdTranslationError, ) - HOUSEHOLD_ID = UUID("00000000-0000-0000-0000-000000000010") diff --git a/tests/unit/services/test_household_service.py b/tests/unit/services/test_household_service.py index 79c2c116e..ee383d97d 100644 --- a/tests/unit/services/test_household_service.py +++ b/tests/unit/services/test_household_service.py @@ -1,6 +1,7 @@ +from uuid import uuid4 + import pytest from sqlalchemy import func, select -from uuid import uuid4 from policyengine_api.data.v1_models import Household, HouseholdMirrorEvent from policyengine_api.services.household_service import ( @@ -15,7 +16,6 @@ valid_request_body, ) - pytest_plugins = ["tests.fixtures.services.household_fixtures"] diff --git a/tests/unit/services/test_local_data_service_boundaries.py b/tests/unit/services/test_local_data_service_boundaries.py index be078a9eb..b88b91c5c 100644 --- a/tests/unit/services/test_local_data_service_boundaries.py +++ b/tests/unit/services/test_local_data_service_boundaries.py @@ -7,7 +7,6 @@ from policyengine_api.services.reform_impacts_service import ReformImpactsService - SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" diff --git a/tests/unit/services/test_orchestration_service_database_boundaries.py b/tests/unit/services/test_orchestration_service_database_boundaries.py index c16e42049..65032050e 100644 --- a/tests/unit/services/test_orchestration_service_database_boundaries.py +++ b/tests/unit/services/test_orchestration_service_database_boundaries.py @@ -4,7 +4,6 @@ import pytest - SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" diff --git a/tests/unit/services/test_orm_service_boundaries.py b/tests/unit/services/test_orm_service_boundaries.py index 1c48c8494..b8d49a12a 100644 --- a/tests/unit/services/test_orm_service_boundaries.py +++ b/tests/unit/services/test_orm_service_boundaries.py @@ -4,7 +4,6 @@ import pytest - SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" diff --git a/tests/unit/services/test_policy_mirroring.py b/tests/unit/services/test_policy_mirroring.py index 1a29aee90..8b9d81e11 100644 --- a/tests/unit/services/test_policy_mirroring.py +++ b/tests/unit/services/test_policy_mirroring.py @@ -2,12 +2,16 @@ from __future__ import annotations -from uuid import UUID from unittest.mock import MagicMock, patch +from uuid import UUID import pytest from sqlalchemy.exc import OperationalError, TimeoutError +from policyengine_api.services.policy_mirroring import ( + PolicyMirrorUnavailableError, + mirror_policy_after_commit, +) from policyengine_api.services.v2.policies.types import ( LegacyPolicyPersistenceResult, LegacyPolicySnapshot, @@ -15,11 +19,6 @@ from policyengine_api.services.v2.policies.validators import ( LegacyPolicyMappingIntegrityError, ) -from policyengine_api.services.policy_mirroring import ( - PolicyMirrorUnavailableError, - mirror_policy_after_commit, -) - POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") diff --git a/tests/unit/services/test_policy_service.py b/tests/unit/services/test_policy_service.py index 1abda22a1..5c916bf8e 100644 --- a/tests/unit/services/test_policy_service.py +++ b/tests/unit/services/test_policy_service.py @@ -6,7 +6,6 @@ from policyengine_api.services.policy_service import PolicyService, PolicySetResult from tests.fixtures.services.policy_service import valid_policy_data - pytest_plugins = ["tests.fixtures.services.policy_service"] diff --git a/tests/unit/services/test_report_output_alias_service.py b/tests/unit/services/test_report_output_alias_service.py index 9377fe06b..f48fd1180 100644 --- a/tests/unit/services/test_report_output_alias_service.py +++ b/tests/unit/services/test_report_output_alias_service.py @@ -8,7 +8,6 @@ ReportOutputAliasService, ) - service = ReportOutputAliasService() diff --git a/tests/unit/services/test_report_output_service.py b/tests/unit/services/test_report_output_service.py index 92344b68e..394cc5973 100644 --- a/tests/unit/services/test_report_output_service.py +++ b/tests/unit/services/test_report_output_service.py @@ -8,7 +8,6 @@ from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.simulation_service import SimulationService - run_service = ReportRunService() SPM_CONFIG = { "forecast_content_sha256": "a" * 64, diff --git a/tests/unit/services/test_report_run_service.py b/tests/unit/services/test_report_run_service.py index 890ef5e5d..565148313 100644 --- a/tests/unit/services/test_report_run_service.py +++ b/tests/unit/services/test_report_run_service.py @@ -3,7 +3,6 @@ from policyengine_api.data.v1_models import ReportOutput, ReportOutputRun from policyengine_api.services.report_run_service import ReportRunService - service = ReportRunService() diff --git a/tests/unit/services/test_report_service_owned_sessions.py b/tests/unit/services/test_report_service_owned_sessions.py index e9103a07b..e855e7a7f 100644 --- a/tests/unit/services/test_report_service_owned_sessions.py +++ b/tests/unit/services/test_report_service_owned_sessions.py @@ -12,7 +12,6 @@ ) from policyengine_api.services.simulation_service import SimulationService - ROUTE_PATH = ( Path(__file__).parents[3] / "policyengine_api" diff --git a/tests/unit/services/test_report_spec_service.py b/tests/unit/services/test_report_spec_service.py index 3db91765e..ad0f4044c 100644 --- a/tests/unit/services/test_report_spec_service.py +++ b/tests/unit/services/test_report_spec_service.py @@ -7,7 +7,6 @@ ReportSpecService, ) - service = ReportSpecService() diff --git a/tests/unit/services/test_run_and_spec_service_database_boundaries.py b/tests/unit/services/test_run_and_spec_service_database_boundaries.py index fcf6cf398..9ace9ddfe 100644 --- a/tests/unit/services/test_run_and_spec_service_database_boundaries.py +++ b/tests/unit/services/test_run_and_spec_service_database_boundaries.py @@ -4,7 +4,6 @@ import pytest - SERVICE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "services" diff --git a/tests/unit/services/test_service_owned_session_architecture.py b/tests/unit/services/test_service_owned_session_architecture.py index 7a1fe24f1..a2829777a 100644 --- a/tests/unit/services/test_service_owned_session_architecture.py +++ b/tests/unit/services/test_service_owned_session_architecture.py @@ -16,7 +16,6 @@ from policyengine_api.services.user_policy_service import UserPolicyService from policyengine_api.services.user_service import UserService - PROJECT_ROOT = Path(__file__).parents[3] PACKAGE_ROOT = PROJECT_ROOT / "policyengine_api" diff --git a/tests/unit/services/test_service_owned_sessions.py b/tests/unit/services/test_service_owned_sessions.py index e2fe338c3..98e76f979 100644 --- a/tests/unit/services/test_service_owned_sessions.py +++ b/tests/unit/services/test_service_owned_sessions.py @@ -13,7 +13,6 @@ from policyengine_api.services.simulation_service import SimulationService from policyengine_api.services.user_service import UserService - ROUTE_ROOT = Path(__file__).parents[3] / "policyengine_api" / "routes" diff --git a/tests/unit/services/test_simulation_run_service.py b/tests/unit/services/test_simulation_run_service.py index 43ba3b80b..f1d606ae0 100644 --- a/tests/unit/services/test_simulation_run_service.py +++ b/tests/unit/services/test_simulation_run_service.py @@ -4,7 +4,6 @@ from policyengine_api.services.simulation_run_service import SimulationRunService from policyengine_api.services.simulation_service import SimulationService - run_service = SimulationRunService() simulation_service = SimulationService() diff --git a/tests/unit/services/test_simulation_spec_service.py b/tests/unit/services/test_simulation_spec_service.py index 26f06d02f..0cf1effd2 100644 --- a/tests/unit/services/test_simulation_spec_service.py +++ b/tests/unit/services/test_simulation_spec_service.py @@ -6,7 +6,6 @@ SimulationSpecService, ) - simulation_service = SimulationService() spec_service = SimulationSpecService() diff --git a/tests/unit/services/test_spm_household_identity.py b/tests/unit/services/test_spm_household_identity.py index 661d3ec03..4a4286ac8 100644 --- a/tests/unit/services/test_spm_household_identity.py +++ b/tests/unit/services/test_spm_household_identity.py @@ -11,7 +11,6 @@ from policyengine_api.services.simulation_service import SimulationService from policyengine_api.spm import SPMValidationError - INPUTS = {"people": {"you": {"age": {"2026": 40}}}} NATIONAL = { "forecast_content_sha256": "a" * 64, diff --git a/tests/unit/services/test_user_policy_mirroring.py b/tests/unit/services/test_user_policy_mirroring.py index cd99ad3c1..ef9f70b89 100644 --- a/tests/unit/services/test_user_policy_mirroring.py +++ b/tests/unit/services/test_user_policy_mirroring.py @@ -2,14 +2,21 @@ from __future__ import annotations -from uuid import UUID from unittest.mock import MagicMock, patch +from uuid import UUID import pytest -from sqlalchemy import event as sqlalchemy_event, select +from sqlalchemy import event as sqlalchemy_event +from sqlalchemy import select from sqlalchemy.exc import OperationalError, TimeoutError from policyengine_api.data.v1_models import UserPolicyMirrorEvent +from policyengine_api.services.user_policy_mirroring import ( + UserPolicyMirrorUnavailableError, + mirror_pending_user_policy_events_after_commit, + mirror_user_policy_after_commit, +) +from policyengine_api.services.user_policy_service import UserPolicyService from policyengine_api.services.v2.policies.types import LegacyPolicySnapshot from policyengine_api.services.v2.user_policies.types import ( LegacyUserPolicyPersistenceResult, @@ -18,13 +25,6 @@ from policyengine_api.services.v2.user_policies.validators import ( LegacyUserPolicyIntegrityError, ) -from policyengine_api.services.user_policy_mirroring import ( - UserPolicyMirrorUnavailableError, - mirror_pending_user_policy_events_after_commit, - mirror_user_policy_after_commit, -) -from policyengine_api.services.user_policy_service import UserPolicyService - ASSOCIATION_ID = UUID("00000000-0000-0000-0000-000000000060") POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") diff --git a/tests/unit/services/test_user_policy_service.py b/tests/unit/services/test_user_policy_service.py index 7fec2014e..f36244427 100644 --- a/tests/unit/services/test_user_policy_service.py +++ b/tests/unit/services/test_user_policy_service.py @@ -10,13 +10,12 @@ IntegrityError, OperationalError, SQLAlchemyError, +) +from sqlalchemy.exc import ( TimeoutError as SQLAlchemyTimeoutError, ) from policyengine_api.data.v1_models import UserPolicy, UserPolicyMirrorEvent -from policyengine_api.services.v2.user_policies.types import ( - LegacyUserPolicyPersistenceResult, -) from policyengine_api.services.user_policy_service import ( UserPolicyCreateResult, UserPolicyMirrorEventIntegrityError, @@ -24,7 +23,9 @@ UserPolicyService, UserPolicyUpdateResult, ) - +from policyengine_api.services.v2.user_policies.types import ( + LegacyUserPolicyPersistenceResult, +) ROUTE_PATH = ( Path(__file__).parents[3] / "policyengine_api" / "routes" / "policy_routes.py" diff --git a/tests/unit/services/test_user_service.py b/tests/unit/services/test_user_service.py index f7804fba5..0f27836b3 100644 --- a/tests/unit/services/test_user_service.py +++ b/tests/unit/services/test_user_service.py @@ -4,7 +4,6 @@ from policyengine_api.services.user_service import UserService from tests.fixtures.services.user_service import valid_user_record - pytest_plugins = ["tests.fixtures.services.user_service"] diff --git a/tests/unit/services/test_worker_spm.py b/tests/unit/services/test_worker_spm.py index b09a614d2..7a1304059 100644 --- a/tests/unit/services/test_worker_spm.py +++ b/tests/unit/services/test_worker_spm.py @@ -274,9 +274,10 @@ def test_uncertified_or_mismatching_worker_is_rejected(capability): ], ) def test_actual_entrypoint_http_preserves_selection_without_data_override(method, path): - import httpx import json + import httpx + requests = [] def handle(request): @@ -376,8 +377,10 @@ def test_explicit_bundle_capability_does_not_inherit_another_bundles_capability( @pytest.mark.parametrize("method", ["run", "run_budget_window_batch"]) def test_canonical_direct_submission_forwards_the_validated_default_bundle(method): - import httpx import json + + import httpx + from policyengine_api.constants import POLICYENGINE_VERSION def handle(request): @@ -442,6 +445,7 @@ def test_typed_worker_poll_errors_reach_api_validation(code, method): def test_worker_cache_receipts_are_required_and_json_roundtrip(): import json + from policyengine_api.worker_spm import validate_worker_result output = { diff --git a/tests/unit/test_alembic_skill.py b/tests/unit/test_alembic_skill.py index d8d7b9f30..c61312390 100644 --- a/tests/unit/test_alembic_skill.py +++ b/tests/unit/test_alembic_skill.py @@ -1,6 +1,5 @@ from policyengine_api.constants import REPO - SKILL = REPO / "docs" / "engineering" / "skills" / "alembic-migrations.md" diff --git a/tests/unit/test_app_engine_decommission.py b/tests/unit/test_app_engine_decommission.py index 9cc17dc48..523199799 100644 --- a/tests/unit/test_app_engine_decommission.py +++ b/tests/unit/test_app_engine_decommission.py @@ -3,7 +3,6 @@ import re from pathlib import Path - REPO = Path(__file__).resolve().parents[2] PUSH_WORKFLOW = REPO / ".github/workflows/push.yml" REMOVED_JOB_NAMES = ( diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index b4660a92e..cb2d1b0d4 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -8,10 +8,11 @@ from unittest.mock import Mock, patch import pytest -import policyengine_api.asgi_factory as asgi_factory from fastapi.testclient import TestClient from flask import Flask, Response, jsonify, make_response, request from flask_cors import CORS + +import policyengine_api.asgi_factory as asgi_factory from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.migration_flags import ( RouteImplementation, diff --git a/tests/unit/test_country.py b/tests/unit/test_country.py index 34561eac7..10e1d57e0 100644 --- a/tests/unit/test_country.py +++ b/tests/unit/test_country.py @@ -1,10 +1,10 @@ -import pytest -import numpy as np -import pandas as pd from datetime import date, datetime from pathlib import Path from types import SimpleNamespace +import numpy as np +import pandas as pd +import pytest from policyengine_core.enums import EnumArray from policyengine_core.model_api import Enum diff --git a/tests/unit/test_country_spm.py b/tests/unit/test_country_spm.py index 937c0e537..3f618b7b0 100644 --- a/tests/unit/test_country_spm.py +++ b/tests/unit/test_country_spm.py @@ -1,7 +1,7 @@ """Exercise the API country adapter, including optional real canonical runtime.""" -from copy import deepcopy import json +from copy import deepcopy from types import SimpleNamespace import numpy as np @@ -11,7 +11,6 @@ from policyengine_api import spm from policyengine_api.country import COUNTRIES, PolicyEngineCountry - # An axes request expands every recognized variable, so a stub system must name # the ones these households carry; anything else (entity membership) is skipped. HOUSEHOLD_VARIABLES = { @@ -397,8 +396,8 @@ def test_real_country_explicit_geography_yields_json_receipt( def real_http_client(real_canonical_country, monkeypatch, orm_session_factory): from flask import Flask - from policyengine_api.extensions import cache from policyengine_api.data.v1_models import Policy + from policyengine_api.extensions import cache from policyengine_api.routes import household_routes from policyengine_api.runtime_cache.core import CacheNamespace from policyengine_api.runtime_cache.fake import InMemoryCacheBackend @@ -653,7 +652,7 @@ def test_uncertifiable_country_receipt_does_not_become_an_internal_failure(monke calculate=lambda variable, period: np.array([123.0]), get_population=lambda entity: SimpleNamespace(get_index=lambda entity_id: 0), spm_config={**selection, "threshold_method": "unreviewed"}, - spm_provenance=lambda: {}, + spm_provenance=dict, ) system = SimpleNamespace( get_variable=lambda name: SimpleNamespace(value_type=float) diff --git a/tests/unit/test_household_event_command.py b/tests/unit/test_household_event_command.py index 9d69871f8..632d75f6d 100644 --- a/tests/unit/test_household_event_command.py +++ b/tests/unit/test_household_event_command.py @@ -13,7 +13,6 @@ ) from scripts import process_v1_household_mirror_event as command - STAGING_ENVIRONMENT = { "DEPLOYMENT_ENVIRONMENT": "staging", "V2_SUPABASE_ENVIRONMENT": "staging", diff --git a/tests/unit/test_observability_collector_verifier.py b/tests/unit/test_observability_collector_verifier.py index 1c3ea1601..fce1d8c49 100644 --- a/tests/unit/test_observability_collector_verifier.py +++ b/tests/unit/test_observability_collector_verifier.py @@ -56,7 +56,7 @@ def __init__(self, _channel): def Export(self, request, *, metadata, timeout): exports[signal] = (request, metadata, timeout) if reject: - raise _Unimplemented() + raise _Unimplemented return _accepted_response() return Stub diff --git a/tests/unit/test_openapi_simulation.py b/tests/unit/test_openapi_simulation.py index 3d9ca1f76..099a3ed43 100644 --- a/tests/unit/test_openapi_simulation.py +++ b/tests/unit/test_openapi_simulation.py @@ -1,13 +1,13 @@ """Served simulation schemas agree with real Flask persistence envelopes.""" -from copy import deepcopy import json +from copy import deepcopy +import jsonschema +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from flask import Flask -import jsonschema -import pytest from policyengine_api.data.v1_models import Simulation from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies diff --git a/tests/unit/test_openapi_spm.py b/tests/unit/test_openapi_spm.py index a7d709a60..dcf9d47b1 100644 --- a/tests/unit/test_openapi_spm.py +++ b/tests/unit/test_openapi_spm.py @@ -4,7 +4,12 @@ from flask import Flask +from policyengine_api.query_parameters import ( + AnnualEconomyQuery, + BudgetWindowEconomyQuery, +) from policyengine_api.routes.system_routes import system_bp +from policyengine_api.specification import load_specification from policyengine_api.spm import ( SPMComparisonProvenance, SPMExecutionProvenance, @@ -12,11 +17,6 @@ SPMRuntimeVersions, SPMSelection, ) -from policyengine_api.query_parameters import ( - AnnualEconomyQuery, - BudgetWindowEconomyQuery, -) -from policyengine_api.specification import load_specification def test_economy_query_and_receipt_schemas_match_public_http_contract(): diff --git a/tests/unit/test_query_parameters.py b/tests/unit/test_query_parameters.py index e7a64dca0..1f7822913 100644 --- a/tests/unit/test_query_parameters.py +++ b/tests/unit/test_query_parameters.py @@ -5,10 +5,10 @@ from typing import Annotated from uuid import UUID, uuid4 +import pytest from fastapi import Depends, FastAPI from fastapi.testclient import TestClient from pydantic import Field, ValidationError -import pytest from werkzeug.datastructures import MultiDict from policyengine_api.fastapi_routes.query_parameters import query_dependency diff --git a/tests/unit/test_request_context.py b/tests/unit/test_request_context.py index 0ccdbd686..2918b2f65 100644 --- a/tests/unit/test_request_context.py +++ b/tests/unit/test_request_context.py @@ -8,7 +8,6 @@ start_observability_id, ) - FIRST_ID = "00000000-0000-4000-8000-000000000001" SECOND_ID = "00000000-0000-4000-8000-000000000002" diff --git a/tests/unit/test_simulation_and_report_routes.py b/tests/unit/test_simulation_and_report_routes.py index 23aa8c76f..559b260d8 100644 --- a/tests/unit/test_simulation_and_report_routes.py +++ b/tests/unit/test_simulation_and_report_routes.py @@ -13,7 +13,6 @@ from policyengine_api.services.report_run_service import ReportRunService from policyengine_api.services.simulation_service import SimulationService - simulation_service = SimulationService() report_service = ReportOutputService() run_service = ReportRunService() diff --git a/tests/unit/test_spm_provenance.py b/tests/unit/test_spm_provenance.py index 08521a1cf..cb579389d 100644 --- a/tests/unit/test_spm_provenance.py +++ b/tests/unit/test_spm_provenance.py @@ -17,7 +17,6 @@ validate_spm_calculation_provenance, ) - SHA256 = "a" * 64 SELECTION = SPMSelection.model_validate( { diff --git a/tests/unit/test_spm_settings.py b/tests/unit/test_spm_settings.py index 6ed564f1f..fa8fa1788 100644 --- a/tests/unit/test_spm_settings.py +++ b/tests/unit/test_spm_settings.py @@ -9,7 +9,6 @@ from policyengine_api import spm - ARTIFACT_HASH = "a" * 64 LEGACY_BUNDLE = { "policyengine_version": "5.2.0", @@ -481,7 +480,7 @@ def test_receipt_reduces_country_diagnostics_to_one_compact_provenance( spm_config={**RESOLVED_SELECTION, "threshold_method": "canonical"}, spm_provenance=lambda: COUNTRY_RECEIPT, ), - SimpleNamespace(spm_config=RESOLVED_SELECTION, spm_provenance=lambda: {}), + SimpleNamespace(spm_config=RESOLVED_SELECTION, spm_provenance=dict), ], ids=["extra-setting", "incomplete-receipt"], ) diff --git a/tests/unit/test_stage6_native_metadata.py b/tests/unit/test_stage6_native_metadata.py index cf28fcf10..c045f3802 100644 --- a/tests/unit/test_stage6_native_metadata.py +++ b/tests/unit/test_stage6_native_metadata.py @@ -5,6 +5,7 @@ import pytest from fastapi.testclient import TestClient from flask import Flask, Response + from policyengine_api.asgi_factory import ( NativeRouteDependencies, create_asgi_app, @@ -15,8 +16,8 @@ RouteImplementationSettings, ) from policyengine_api.request_context import REQUEST_ID_HEADER -from policyengine_api.utils.streaming_json import STREAMING_THRESHOLD_BYTES from policyengine_api.utils.payload_validators import validate_country +from policyengine_api.utils.streaming_json import STREAMING_THRESHOLD_BYTES class _HealthySimulationGateway: @@ -50,7 +51,7 @@ def _dependencies(reader: _MetadataReader) -> NativeRouteDependencies: readiness_probe=lambda: True, gateway_client_factory=_HealthySimulationGateway, metadata_reader_factory=lambda: reader, - specification_provider=lambda: {}, + specification_provider=dict, ) diff --git a/tests/unit/test_worker_spm_release.py b/tests/unit/test_worker_spm_release.py index 9df37be59..b3a043ecc 100644 --- a/tests/unit/test_worker_spm_release.py +++ b/tests/unit/test_worker_spm_release.py @@ -7,7 +7,6 @@ from policyengine_api import constants, worker_spm, worker_spm_release from policyengine_api.spm import SPMValidationError - FIXTURE_BUNDLE_VERSION = "fixture-policyengine" FIXTURE_PACKAGE_VERSIONS = { "policyengine": FIXTURE_BUNDLE_VERSION, diff --git a/tests/unit/test_yearly_var_removal.py b/tests/unit/test_yearly_var_removal.py index 9c13fd378..c2dc439fd 100644 --- a/tests/unit/test_yearly_var_removal.py +++ b/tests/unit/test_yearly_var_removal.py @@ -5,7 +5,6 @@ add_yearly_variables, ) - TEST_YEAR = "2023" diff --git a/tests/unit/utils/test_input_validation.py b/tests/unit/utils/test_input_validation.py index 9cd70efa2..922e22ac1 100644 --- a/tests/unit/utils/test_input_validation.py +++ b/tests/unit/utils/test_input_validation.py @@ -1,6 +1,5 @@ from policyengine_api.utils.input_validation import find_unrecognized_inputs - METADATA = { "variables": { "age": {"entity": "person"}, diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 0bc4eb7bb..95a578b55 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -1,27 +1,27 @@ """Unit guards for the isolated and qualified v2 Alembic environment.""" -from io import StringIO -from pathlib import Path import re import shutil +from io import StringIO +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock +import pytest from alembic import command from alembic.config import Config from alembic.script import ScriptDirectory from alembic.script.revision import ResolutionError from alembic.util import CommandError -import pytest from policyengine_api.constants import REPO from policyengine_api.data.v1_models import V1Base from policyengine_api.data.v2.migration_target import ( - ConfiguredSupabaseTarget, DISPOSABLE_DATABASE_NAME, MIGRATION_ROLE, V2_ALEMBIC_DISPOSABLE_TEST, V2_ALEMBIC_HEAD_REVISION, + ConfiguredSupabaseTarget, V2AlembicSettings, V2MigrationTargetError, load_v2_alembic_settings, diff --git a/tests/unit/v2/test_catalog_extraction.py b/tests/unit/v2/test_catalog_extraction.py index cb43c4508..e060c1d43 100644 --- a/tests/unit/v2/test_catalog_extraction.py +++ b/tests/unit/v2/test_catalog_extraction.py @@ -3,9 +3,9 @@ from __future__ import annotations import ast -from datetime import datetime, timezone -from importlib.metadata import PackageNotFoundError import json +from datetime import UTC, datetime +from importlib.metadata import PackageNotFoundError from pathlib import Path from types import SimpleNamespace @@ -108,7 +108,7 @@ def item(self): assert normalize_json_value( { - "date": datetime(2026, 1, 1, tzinfo=timezone.utc), + "date": datetime(2026, 1, 1, tzinfo=UTC), "values": (Scalar(), True, 1.5), } ) == { @@ -323,7 +323,7 @@ def test_rejects_duplicate_keys_invalid_intervals_and_unknown_region_types() -> invalid_interval = source_models() invalid_interval["us"].parameters_by_name["gov.example.rate"].parameter_values[ 0 - ].end_date = datetime(2026, 1, 1, tzinfo=timezone.utc) + ].end_date = datetime(2026, 1, 1, tzinfo=UTC) with pytest.raises(CatalogExtractionError, match="canonical inclusive intervals"): extract_catalog( bundle=bundle(), @@ -367,8 +367,8 @@ def test_collapses_equal_parameter_values_at_the_same_effective_date() -> None: values = catalog.country("us").parameters[0].values assert [(value.start_date, value.value_json) for value in values] == [ - (datetime(2025, 1, 1, tzinfo=timezone.utc), 0.1), - (datetime(2026, 1, 1, tzinfo=timezone.utc), 0.2), + (datetime(2025, 1, 1, tzinfo=UTC), 0.1), + (datetime(2026, 1, 1, tzinfo=UTC), 0.2), ] @@ -420,12 +420,12 @@ def test_preserves_consecutive_day_parameter_value_intervals() -> None: parameter.parameter_values = [ SimpleNamespace( value=0.1, - start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), - end_date=datetime(2026, 1, 1, tzinfo=timezone.utc), + start_date=datetime(2026, 1, 1, tzinfo=UTC), + end_date=datetime(2026, 1, 1, tzinfo=UTC), ), SimpleNamespace( value=0.2, - start_date=datetime(2026, 1, 2, tzinfo=timezone.utc), + start_date=datetime(2026, 1, 2, tzinfo=UTC), end_date=None, ), ] @@ -439,7 +439,7 @@ def test_preserves_consecutive_day_parameter_value_intervals() -> None: values = catalog.country("us").parameters[0].values assert values[0].start_date == values[0].end_date - assert values[1].start_date == datetime(2026, 1, 2, tzinfo=timezone.utc) + assert values[1].start_date == datetime(2026, 1, 2, tzinfo=UTC) assert values[1].end_date is None diff --git a/tests/unit/v2/test_catalog_initialization.py b/tests/unit/v2/test_catalog_initialization.py index 43211aada..6f7dc38ae 100644 --- a/tests/unit/v2/test_catalog_initialization.py +++ b/tests/unit/v2/test_catalog_initialization.py @@ -20,7 +20,6 @@ normalized_catalog, ) - ENVIRONMENT = { V2_DATA_WRITE_DATABASE_URL: ( "postgresql+psycopg://data-writer:test-password@db." diff --git a/tests/unit/v2/test_catalog_publication.py b/tests/unit/v2/test_catalog_publication.py index 187101086..37346a899 100644 --- a/tests/unit/v2/test_catalog_publication.py +++ b/tests/unit/v2/test_catalog_publication.py @@ -6,10 +6,10 @@ from pathlib import Path from unittest.mock import patch -from alembic.config import Config -from alembic.script import ScriptDirectory import pytest import sqlalchemy as sa +from alembic.config import Config +from alembic.script import ScriptDirectory from sqlalchemy.dialects import postgresql from sqlalchemy.schema import CreateTable from sqlalchemy.sql.elements import TextClause @@ -32,7 +32,6 @@ ) from tests.fixtures.v2_catalog import normalized_catalog - REPO = Path(__file__).parents[3] diff --git a/tests/unit/v2/test_data_crud_boundaries.py b/tests/unit/v2/test_data_crud_boundaries.py index f7d03f0f6..935cf06a4 100644 --- a/tests/unit/v2/test_data_crud_boundaries.py +++ b/tests/unit/v2/test_data_crud_boundaries.py @@ -7,7 +7,6 @@ import pytest - PROJECT_ROOT = Path(__file__).parents[3] SERVICES_ROOT = PROJECT_ROOT / "policyengine_api" / "services" / "v2" @@ -41,9 +40,8 @@ def _imported_modules(path: Path) -> set[str]: def _has_database_query_dependency(path: Path) -> bool: return any( module == "sqlalchemy" - or module.startswith("sqlalchemy.") + or module.startswith(("sqlalchemy.", "sqlmodel.")) or module == "sqlmodel" - or module.startswith("sqlmodel.") for module in _imported_modules(path) ) diff --git a/tests/unit/v2/test_household_inputs.py b/tests/unit/v2/test_household_inputs.py index 36a7b1a08..3cb3a2c39 100644 --- a/tests/unit/v2/test_household_inputs.py +++ b/tests/unit/v2/test_household_inputs.py @@ -2,8 +2,8 @@ from __future__ import annotations -from pydantic import ValidationError import pytest +from pydantic import ValidationError from policyengine_api.services.v2.households.services import normalize_creation_input from policyengine_api.services.v2.households.types import ( diff --git a/tests/unit/v2/test_household_migration_qualification.py b/tests/unit/v2/test_household_migration_qualification.py index 2ae402d53..a264d61ba 100644 --- a/tests/unit/v2/test_household_migration_qualification.py +++ b/tests/unit/v2/test_household_migration_qualification.py @@ -14,7 +14,6 @@ V2_SUPABASE_PROJECT_REF, ) - ENVIRONMENT = { V2_MIGRATION_DATABASE_URL: ( "postgresql+psycopg://migration:test-password@db." diff --git a/tests/unit/v2/test_household_models.py b/tests/unit/v2/test_household_models.py index 4b1025f71..f0798e35e 100644 --- a/tests/unit/v2/test_household_models.py +++ b/tests/unit/v2/test_household_models.py @@ -12,11 +12,11 @@ from policyengine_api.data.v1_models import HouseholdMirrorEvent from policyengine_api.data.v2.models import ( + V2_METADATA, Household, LegacyHouseholdMapping, User, UserHouseholdAssociation, - V2_METADATA, ) diff --git a/tests/unit/v2/test_household_persistence_statements.py b/tests/unit/v2/test_household_persistence_statements.py index 731bbaaef..735ec9930 100644 --- a/tests/unit/v2/test_household_persistence_statements.py +++ b/tests/unit/v2/test_household_persistence_statements.py @@ -6,8 +6,8 @@ from unittest.mock import Mock from uuid import uuid4 -from sqlalchemy.dialects import postgresql import pytest +from sqlalchemy.dialects import postgresql from policyengine_api.data.v2.models import Household, LegacyHouseholdMapping from policyengine_api.services.v2.households.database_connectors import creates diff --git a/tests/unit/v2/test_household_query.py b/tests/unit/v2/test_household_query.py index d11016f60..ccfa5d730 100644 --- a/tests/unit/v2/test_household_query.py +++ b/tests/unit/v2/test_household_query.py @@ -2,13 +2,13 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID -from sqlmodel import Session, create_engine import pytest +from sqlmodel import Session, create_engine -from policyengine_api.data.v2.models import Household, V2_METADATA +from policyengine_api.data.v2.models import V2_METADATA, Household from policyengine_api.services.v2.households.services import ( read_complete_household, read_household_page, @@ -19,7 +19,7 @@ def _stored_households(): engine = create_engine("sqlite://") V2_METADATA.create_all(engine) - created = datetime(2026, 1, 1, tzinfo=timezone.utc) + created = datetime(2026, 1, 1, tzinfo=UTC) first_id = UUID("00000000-0000-0000-0000-000000000010") second_id = UUID("00000000-0000-0000-0000-000000000020") other_id = UUID("00000000-0000-0000-0000-000000000030") diff --git a/tests/unit/v2/test_household_routes.py b/tests/unit/v2/test_household_routes.py index df8e3f7ec..64c0e1343 100644 --- a/tests/unit/v2/test_household_routes.py +++ b/tests/unit/v2/test_household_routes.py @@ -2,12 +2,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 +import pytest from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from sqlalchemy.exc import OperationalError, TimeoutError from policyengine_api.asgi_factory import create_asgi_app @@ -29,9 +29,8 @@ HouseholdValidationError, ) - HOUSEHOLD_ID = UUID("00000000-0000-0000-0000-000000000010") -CREATED_AT = datetime(2026, 1, 1, tzinfo=timezone.utc) +CREATED_AT = datetime(2026, 1, 1, tzinfo=UTC) def _document() -> dict[str, object]: @@ -120,7 +119,7 @@ def fallback(resource: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_household_service_factory=lambda: service, ) settings = RouteImplementationSettings( diff --git a/tests/unit/v2/test_household_translation.py b/tests/unit/v2/test_household_translation.py index 9b12fd92a..53b5d141d 100644 --- a/tests/unit/v2/test_household_translation.py +++ b/tests/unit/v2/test_household_translation.py @@ -5,8 +5,8 @@ import json from pathlib import Path -from pydantic import ValidationError import pytest +from pydantic import ValidationError from policyengine_api.services.v2.households.transformations import ( LEGACY_COLLECTION_NAMES, @@ -18,7 +18,6 @@ LegacyHouseholdTranslationError, ) - DATA = Path(__file__).parents[2] / "data" @@ -189,6 +188,7 @@ def test_saved_spm_translation_preserves_selection_without_current_bundle_resolu ) -> None: from copy import deepcopy from unittest.mock import Mock + from policyengine_api import spm from policyengine_api.services.v2.households.transformations import ( canonicalize_household, diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index 5de6c3efb..768cf1843 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -1,9 +1,9 @@ """Import/startup tests for dormant and fail-closed v2 configuration.""" import os -from pathlib import Path import subprocess import sys +from pathlib import Path import pytest @@ -15,7 +15,6 @@ load_v2_runtime_database_settings, ) - V2_ENVIRONMENT_NAMES = ( "V2_RUNTIME_DATABASE_URL", "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE", diff --git a/tests/unit/v2/test_metadata_deployment.py b/tests/unit/v2/test_metadata_deployment.py index c4619c341..c636c347d 100644 --- a/tests/unit/v2/test_metadata_deployment.py +++ b/tests/unit/v2/test_metadata_deployment.py @@ -2,10 +2,9 @@ from __future__ import annotations -from pathlib import Path import re import subprocess - +from pathlib import Path REPO = Path(__file__).resolve().parents[3] diff --git a/tests/unit/v2/test_metadata_routes.py b/tests/unit/v2/test_metadata_routes.py index 480d4a0c4..229b5e67c 100644 --- a/tests/unit/v2/test_metadata_routes.py +++ b/tests/unit/v2/test_metadata_routes.py @@ -2,13 +2,13 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from unittest.mock import patch from uuid import uuid4 +import pytest from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.data.v2.catalog.catalog_selection import ( @@ -17,6 +17,13 @@ MetadataCatalogVersionNotFoundError, UnsupportedPreviewCountryError, ) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes import dependencies as route_dependencies +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) from policyengine_api.services.v2.metadata.types import ( MetadataCanonicalParameterValue, MetadataDataset, @@ -38,13 +45,6 @@ InvalidMetadataPageError, MetadataResourceNotFoundError, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.fastapi_routes import dependencies as route_dependencies -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.migration_flags import ( - RouteImplementation, - RouteImplementationSettings, -) class ResourceReader: @@ -125,7 +125,7 @@ def _resource_results() -> tuple[dict[str, object], dict[str, object]]: id=parameter_value_id, parameter_id=parameter_id, value=0.1, - start_date=datetime(2025, 1, 1, tzinfo=timezone.utc), + start_date=datetime(2025, 1, 1, tzinfo=UTC), end_date=None, ) dataset = MetadataDataset( @@ -238,7 +238,7 @@ def v1_metadata(country_id: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_metadata_reader_factory=factory, ) settings = RouteImplementationSettings( diff --git a/tests/unit/v2/test_metadata_service.py b/tests/unit/v2/test_metadata_service.py index fe1069cb5..9880f265e 100644 --- a/tests/unit/v2/test_metadata_service.py +++ b/tests/unit/v2/test_metadata_service.py @@ -3,7 +3,7 @@ from __future__ import annotations import ast -from datetime import datetime, timezone +from datetime import UTC, datetime from pathlib import Path from uuid import uuid4 @@ -18,15 +18,8 @@ MetadataCatalogVersionNotFoundError, UnsupportedPreviewCountryError, ) -from policyengine_api.services.v2.metadata.database_session import ( - MetadataDatabaseSession, -) -from policyengine_api.services.v2.metadata.services import V2MetadataService -from policyengine_api.services.v2.metadata.validators import ( - InvalidMetadataPageError, - MetadataResourceNotFoundError, -) from policyengine_api.data.v2.models import ( + V2_METADATA, Dataset, Parameter, ParameterNode, @@ -35,9 +28,16 @@ RegionType, TaxBenefitModel, TaxBenefitModelVersion, - V2_METADATA, Variable, ) +from policyengine_api.services.v2.metadata.database_session import ( + MetadataDatabaseSession, +) +from policyengine_api.services.v2.metadata.services import V2MetadataService +from policyengine_api.services.v2.metadata.validators import ( + InvalidMetadataPageError, + MetadataResourceNotFoundError, +) from tests.fixtures.v2_catalog import POLICYENGINE_VERSION, normalized_catalog @@ -320,7 +320,7 @@ def test_parameter_values_are_separate_canonical_resources( id=uuid4(), parameter_id=parameter.id, value_json=0.9, - start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), + start_date=datetime(2026, 1, 1, tzinfo=UTC), end_date=None, policy_id=uuid4(), ) @@ -335,7 +335,7 @@ def test_parameter_values_are_separate_canonical_resources( "us", parameter_id=parameter.id, current=True, - now=datetime(2026, 6, 1, tzinfo=timezone.utc), + now=datetime(2026, 6, 1, tzinfo=UTC), ) assert [item.value for item in all_values.items] == [0.2, 0.1] @@ -346,9 +346,9 @@ def test_parameter_values_are_separate_canonical_resources( @pytest.mark.parametrize( ("selected_time", "expected_value"), [ - (datetime(2025, 12, 31, tzinfo=timezone.utc), 0.1), - (datetime(2025, 12, 31, 23, 59, tzinfo=timezone.utc), 0.1), - (datetime(2026, 1, 1, tzinfo=timezone.utc), 0.2), + (datetime(2025, 12, 31, tzinfo=UTC), 0.1), + (datetime(2025, 12, 31, 23, 59, tzinfo=UTC), 0.1), + (datetime(2026, 1, 1, tzinfo=UTC), 0.2), ], ) def test_current_parameter_value_uses_inclusive_effective_dates( diff --git a/tests/unit/v2/test_model_persistence.py b/tests/unit/v2/test_model_persistence.py index 2a4cc854e..95a790791 100644 --- a/tests/unit/v2/test_model_persistence.py +++ b/tests/unit/v2/test_model_persistence.py @@ -1,6 +1,6 @@ """Canonical SQLModel persistence and bounded SQLAlchemy escape-hatch tests.""" -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from pathlib import Path import pytest @@ -9,6 +9,7 @@ from sqlmodel import Session, create_engine, select from policyengine_api.data.v2.models import ( + V2_METADATA, Dynamic, LegacyPolicyMapping, LegacyUserMapping, @@ -20,7 +21,6 @@ TaxBenefitModelVersion, User, UserPolicy, - V2_METADATA, ) from policyengine_api.data.v2.models.base import DIRECT_SQLALCHEMY_EXCEPTIONS @@ -140,7 +140,7 @@ def test_canonical_parameter_values_are_unique_by_parameter_and_start_date() -> name="gov.example.rate", tax_benefit_model_version=version, ) - start_date = datetime(2026, 1, 1, tzinfo=timezone.utc) + start_date = datetime(2026, 1, 1, tzinfo=UTC) with Session(engine) as session: session.add_all( @@ -186,7 +186,7 @@ def test_policy_parameter_value_owner_period_and_identity_constraints() -> None: engine = _relational_sqlite_engine() _model, _version, parameter, policy = _policy_graph(content_hash="b" * 64) dynamic = Dynamic(name="dynamic") - start = datetime(2026, 1, 1, tzinfo=timezone.utc) + start = datetime(2026, 1, 1, tzinfo=UTC) with Session(engine) as session: session.add_all([parameter, policy, dynamic]) diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index bf9b3e2f4..f21825af0 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -7,6 +7,7 @@ from policyengine_api.data.v1_models import V1Base from policyengine_api.data.v2.models import ( + V2_METADATA, DatasetVersion, Dynamic, Household, @@ -24,10 +25,8 @@ UserPolicy, UserReportAssociation, UserSimulationAssociation, - V2_METADATA, ) - RUN_OUTPUT_TABLES = frozenset( { "aggregates", diff --git a/tests/unit/v2/test_policy_canonicalization.py b/tests/unit/v2/test_policy_canonicalization.py index 8d0e9ffcd..98c7871b3 100644 --- a/tests/unit/v2/test_policy_canonicalization.py +++ b/tests/unit/v2/test_policy_canonicalization.py @@ -2,9 +2,9 @@ from __future__ import annotations -from datetime import datetime, timezone import hashlib import math +from datetime import UTC, datetime from uuid import UUID, uuid4 import pytest @@ -18,7 +18,6 @@ ResolvedPolicyCreationInput, ) - MODEL_ID = UUID("00000000-0000-0000-0000-000000000010") MODEL_VERSION_ID = UUID("00000000-0000-0000-0000-000000000020") FIRST_PARAMETER_ID = UUID("00000000-0000-0000-0000-000000000030") @@ -220,7 +219,7 @@ def test_datetime_objects_are_rendered_at_fixed_microsecond_precision() -> None: 0, 0, 42, - tzinfo=timezone.utc, + tzinfo=UTC, ), } ] diff --git a/tests/unit/v2/test_policy_catalog.py b/tests/unit/v2/test_policy_catalog.py index 3133026dd..3dc5f5235 100644 --- a/tests/unit/v2/test_policy_catalog.py +++ b/tests/unit/v2/test_policy_catalog.py @@ -4,17 +4,17 @@ from uuid import uuid4 -from sqlmodel import Session, create_engine import pytest +from sqlmodel import Session, create_engine from policyengine_api.data.v2.catalog.catalog_selection import ( MetadataCatalogVersionNotFoundError, ) from policyengine_api.data.v2.models import ( + V2_METADATA, Parameter, TaxBenefitModel, TaxBenefitModelVersion, - V2_METADATA, ) from policyengine_api.services.v2.policies.services import ( resolve_policy_creation_input, diff --git a/tests/unit/v2/test_policy_inputs.py b/tests/unit/v2/test_policy_inputs.py index 9b3153604..aa07949a4 100644 --- a/tests/unit/v2/test_policy_inputs.py +++ b/tests/unit/v2/test_policy_inputs.py @@ -2,12 +2,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from uuid import uuid4 -from pydantic import ValidationError import pytest +from pydantic import ValidationError from policyengine_api.services.v2.policies.types import ( NativePolicyCreationInput, @@ -49,7 +49,7 @@ def test_input_normalizes_country_and_effective_dates_to_utc() -> None: 2026, 1, 1, - tzinfo=timezone.utc, + tzinfo=UTC, ) @@ -139,8 +139,7 @@ def test_parameter_value_count_is_bounded_but_empty_policy_is_valid() -> None: { **repeated, "parameter_id": uuid4(), - "start_date": datetime(2026, 1, 1, tzinfo=timezone.utc) - + timedelta(days=index), + "start_date": datetime(2026, 1, 1, tzinfo=UTC) + timedelta(days=index), } for index in range(MAXIMUM_POLICY_PARAMETER_VALUES + 1) ] diff --git a/tests/unit/v2/test_policy_legacy_translation.py b/tests/unit/v2/test_policy_legacy_translation.py index 8bb58e663..7f5d3f236 100644 --- a/tests/unit/v2/test_policy_legacy_translation.py +++ b/tests/unit/v2/test_policy_legacy_translation.py @@ -2,20 +2,25 @@ from __future__ import annotations -from datetime import datetime, timezone import json +from datetime import UTC, datetime +import pytest from pydantic import ValidationError from sqlalchemy.orm import sessionmaker from sqlmodel import Session, create_engine -import pytest from policyengine_api.data.v2.models import ( + V2_METADATA, Parameter, TaxBenefitModel, TaxBenefitModelVersion, - V2_METADATA, ) +from policyengine_api.services.policy_mirroring import ( + PolicyMirrorUnavailableError, + mirror_policy_after_commit, +) +from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.v2.policies.database_connectors.reads import ( read_parameters_by_name, read_policy_catalog, @@ -29,11 +34,6 @@ from policyengine_api.services.v2.policies.validators import ( LegacyPolicyTranslationError, ) -from policyengine_api.services.policy_service import PolicyService -from policyengine_api.services.policy_mirroring import ( - PolicyMirrorUnavailableError, - mirror_policy_after_commit, -) from tests.fixtures.local_v1_database import create_test_v1_schema @@ -104,16 +104,16 @@ def _translate( def test_year_day_and_explicit_range_periods_are_inclusive_utc() -> None: assert parse_legacy_period("2026") == ( - datetime(2026, 1, 1, tzinfo=timezone.utc), - datetime(2026, 12, 31, tzinfo=timezone.utc), + datetime(2026, 1, 1, tzinfo=UTC), + datetime(2026, 12, 31, tzinfo=UTC), ) assert parse_legacy_period("2026-03-02") == ( - datetime(2026, 3, 2, tzinfo=timezone.utc), - datetime(2026, 3, 2, tzinfo=timezone.utc), + datetime(2026, 3, 2, tzinfo=UTC), + datetime(2026, 3, 2, tzinfo=UTC), ) assert parse_legacy_period("2026-02-01.2026-02-28") == ( - datetime(2026, 2, 1, tzinfo=timezone.utc), - datetime(2026, 2, 28, tzinfo=timezone.utc), + datetime(2026, 2, 1, tzinfo=UTC), + datetime(2026, 2, 28, tzinfo=UTC), ) diff --git a/tests/unit/v2/test_policy_migration_qualification.py b/tests/unit/v2/test_policy_migration_qualification.py index cc3f63201..8a3bfde66 100644 --- a/tests/unit/v2/test_policy_migration_qualification.py +++ b/tests/unit/v2/test_policy_migration_qualification.py @@ -14,7 +14,6 @@ V2_SUPABASE_PROJECT_REF, ) - ENVIRONMENT = { V2_MIGRATION_DATABASE_URL: ( "postgresql+psycopg://migration:test-password@db." diff --git a/tests/unit/v2/test_policy_query.py b/tests/unit/v2/test_policy_query.py index ebc9748cd..4fa5b129b 100644 --- a/tests/unit/v2/test_policy_query.py +++ b/tests/unit/v2/test_policy_query.py @@ -2,19 +2,19 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID -from sqlmodel import Session, create_engine import pytest +from sqlmodel import Session, create_engine from policyengine_api.data.v2.models import ( + V2_METADATA, Parameter, ParameterValue, Policy, TaxBenefitModel, TaxBenefitModelVersion, - V2_METADATA, ) from policyengine_api.services.v2.policies.services import ( read_complete_policy, @@ -38,7 +38,7 @@ def _stored_policies(): ) alpha = Parameter(name="gov.alpha", tax_benefit_model_version=version) zeta = Parameter(name="gov.zeta", tax_benefit_model_version=version) - created = datetime(2026, 1, 1, tzinfo=timezone.utc) + created = datetime(2026, 1, 1, tzinfo=UTC) first = Policy( id=UUID("00000000-0000-0000-0000-000000000010"), country_id="us", diff --git a/tests/unit/v2/test_policy_routes.py b/tests/unit/v2/test_policy_routes.py index 900de1abd..e8385d593 100644 --- a/tests/unit/v2/test_policy_routes.py +++ b/tests/unit/v2/test_policy_routes.py @@ -2,12 +2,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 +import pytest from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from sqlalchemy.exc import OperationalError, TimeoutError from policyengine_api.asgi_factory import create_asgi_app @@ -15,6 +15,12 @@ MetadataCatalogUnavailableError, MetadataCatalogVersionNotFoundError, ) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) from policyengine_api.services.v2.policies.types import ( NativePolicyCreation, PolicyPage, @@ -27,20 +33,13 @@ PolicyCreationIntegrityError, PolicyNotFoundError, ) -from policyengine_api.data.v2.settings import V2ConfigurationError -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.migration_flags import ( - RouteImplementation, - RouteImplementationSettings, -) - POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") MODEL_ID = UUID("00000000-0000-0000-0000-000000000020") MODEL_VERSION_ID = UUID("00000000-0000-0000-0000-000000000030") PARAMETER_ID = UUID("00000000-0000-0000-0000-000000000040") VALUE_ID = UUID("00000000-0000-0000-0000-000000000050") -CREATED_AT = datetime(2026, 1, 1, tzinfo=timezone.utc) +CREATED_AT = datetime(2026, 1, 1, tzinfo=UTC) def _policy_read(*, policy_id: UUID = POLICY_ID, country_id: str = "us") -> PolicyRead: @@ -110,7 +109,7 @@ def fallback(resource: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_policy_service_factory=lambda: service, ) settings = RouteImplementationSettings( diff --git a/tests/unit/v2/test_report_runs.py b/tests/unit/v2/test_report_runs.py index a203e3c15..d40ca07b2 100644 --- a/tests/unit/v2/test_report_runs.py +++ b/tests/unit/v2/test_report_runs.py @@ -1,9 +1,9 @@ """Report definition, rerun, idempotency, worker, and selector tests.""" +import threading from concurrent.futures import ThreadPoolExecutor -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from pathlib import Path -import threading from uuid import UUID, uuid4 import pytest @@ -12,6 +12,7 @@ from sqlmodel import Session, create_engine, select from policyengine_api.data.v2.models import ( + V2_METADATA, AggregateOutput, AggregateType, Dataset, @@ -23,7 +24,6 @@ SimulationType, TaxBenefitModel, TaxBenefitModelVersion, - V2_METADATA, ) from policyengine_api.data.v2.report_runs import ( ReportRunStateError, @@ -36,10 +36,9 @@ set_report_type, ) - DEPLOYED_COUNTRY_VERSIONS = {"us": "1.2.3"} DEPLOYED_POLICYENGINE_VERSION = "4.5.6" -NOW = datetime(2026, 8, 14, 12, 0, tzinfo=timezone.utc) +NOW = datetime(2026, 8, 14, 12, 0, tzinfo=UTC) @pytest.fixture diff --git a/tests/unit/v2/test_settings.py b/tests/unit/v2/test_settings.py index f093fd275..9baccde13 100644 --- a/tests/unit/v2/test_settings.py +++ b/tests/unit/v2/test_settings.py @@ -15,7 +15,6 @@ load_v2_runtime_database_settings, ) - PROJECT_REF = "abcdefghijklmnopqrst" TARGET_ENVIRONMENT = { V2_SUPABASE_PROJECT_REF: PROJECT_REF, diff --git a/tests/unit/v2/test_stage12_comparison_models.py b/tests/unit/v2/test_stage12_comparison_models.py index 21c264981..91c7fd1e9 100644 --- a/tests/unit/v2/test_stage12_comparison_models.py +++ b/tests/unit/v2/test_stage12_comparison_models.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID import sqlalchemy as sa @@ -10,13 +10,13 @@ from sqlmodel import Session, create_engine from policyengine_api.data.v2.models import ( + V2_METADATA, Stage12AggregationStatus, Stage12ComparisonReport, Stage12ComparisonSimulation, Stage12ResultComparisonStatus, Stage12RunStatus, Stage12SimulationRole, - V2_METADATA, ) from policyengine_api.services.v2.comparison_runs.types import ( ComparisonRunAggregationStatus, @@ -25,7 +25,7 @@ ) from policyengine_api.services.v2.simulations.types import SimulationRole -NOW = datetime(2026, 9, 14, tzinfo=timezone.utc) +NOW = datetime(2026, 9, 14, tzinfo=UTC) EVALUATION_ID = UUID("00000000-0000-4000-8000-000000000001") DIGEST_A = "a" * 64 DIGEST_B = "b" * 64 diff --git a/tests/unit/v2/test_stage12_comparison_service.py b/tests/unit/v2/test_stage12_comparison_service.py index 688d7412f..072b655f0 100644 --- a/tests/unit/v2/test_stage12_comparison_service.py +++ b/tests/unit/v2/test_stage12_comparison_service.py @@ -2,12 +2,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from unittest.mock import Mock from uuid import UUID, uuid4 -from sqlalchemy.dialects import postgresql import pytest +from sqlalchemy.dialects import postgresql from policyengine_api.data.v2.models import ( Stage12ComparisonReport, @@ -37,7 +37,7 @@ require_result_comparison_transition, ) -NOW = datetime(2026, 9, 14, tzinfo=timezone.utc) +NOW = datetime(2026, 9, 14, tzinfo=UTC) EVALUATION_ID = UUID("00000000-0000-4000-8000-000000000001") SIMULATION_ID = UUID("00000000-0000-4000-8000-000000000002") DIGEST_A = "a" * 64 diff --git a/tests/unit/v2/test_stage12_worker_contracts.py b/tests/unit/v2/test_stage12_worker_contracts.py index 330af0483..75229fc3a 100644 --- a/tests/unit/v2/test_stage12_worker_contracts.py +++ b/tests/unit/v2/test_stage12_worker_contracts.py @@ -5,8 +5,8 @@ from copy import deepcopy from uuid import UUID -from pydantic import ValidationError import pytest +from pydantic import ValidationError from policyengine_api.services.v2.reports.types import ( AggregateReportArtifactDescriptor, diff --git a/tests/unit/v2/test_user_household_routes.py b/tests/unit/v2/test_user_household_routes.py index 4be95201f..cbd179b7c 100644 --- a/tests/unit/v2/test_user_household_routes.py +++ b/tests/unit/v2/test_user_household_routes.py @@ -2,12 +2,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 +import pytest from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from sqlalchemy.exc import IntegrityError, OperationalError, TimeoutError from policyengine_api.asgi_factory import create_asgi_app @@ -28,12 +28,11 @@ UserHouseholdNotFoundError, ) - ASSOCIATION_ID = UUID("00000000-0000-0000-0000-000000000060") HOUSEHOLD_ID = UUID("00000000-0000-0000-0000-000000000010") REPLACEMENT_ID = UUID("00000000-0000-0000-0000-000000000011") USER_ID = UUID("00000000-0000-0000-0000-000000000070") -CREATED_AT = datetime(2026, 1, 1, tzinfo=timezone.utc) +CREATED_AT = datetime(2026, 1, 1, tzinfo=UTC) def _association_read( @@ -123,7 +122,7 @@ def fallback(resource: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_user_household_service_factory=lambda: service, ) settings = RouteImplementationSettings( diff --git a/tests/unit/v2/test_user_household_service.py b/tests/unit/v2/test_user_household_service.py index bd6c59d78..6a4a2bab5 100644 --- a/tests/unit/v2/test_user_household_service.py +++ b/tests/unit/v2/test_user_household_service.py @@ -10,10 +10,10 @@ from sqlmodel import Session, create_engine, select from policyengine_api.data.v2.models import ( + V2_METADATA, Household, User, UserHouseholdAssociation, - V2_METADATA, ) from policyengine_api.services.v2.user_households.database_session import ( UserHouseholdDatabaseSession, @@ -32,7 +32,6 @@ UserHouseholdNotFoundError, ) - USER_ID = UUID("00000000-0000-0000-0000-000000000070") OTHER_USER_ID = UUID("00000000-0000-0000-0000-000000000071") diff --git a/tests/unit/v2/test_user_policy_legacy.py b/tests/unit/v2/test_user_policy_legacy.py index 621ff217e..12e594c0d 100644 --- a/tests/unit/v2/test_user_policy_legacy.py +++ b/tests/unit/v2/test_user_policy_legacy.py @@ -20,7 +20,6 @@ LegacyUserPolicyIntegrityError, ) - POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") USER_ID = UUID("00000000-0000-0000-0000-000000000070") diff --git a/tests/unit/v2/test_user_policy_routes.py b/tests/unit/v2/test_user_policy_routes.py index 922a99e08..3bd62949b 100644 --- a/tests/unit/v2/test_user_policy_routes.py +++ b/tests/unit/v2/test_user_policy_routes.py @@ -2,16 +2,21 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 +import pytest from fastapi.testclient import TestClient from flask import Flask, jsonify -import pytest from sqlalchemy.exc import OperationalError, SQLAlchemyError, TimeoutError from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) from policyengine_api.services.v2.user_policies.types import ( UserPolicyPage, UserPolicyRead, @@ -22,17 +27,11 @@ AssociationUserNotFoundError, UserPolicyNotFoundError, ) -from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.migration_flags import ( - RouteImplementation, - RouteImplementationSettings, -) - ASSOCIATION_ID = UUID("00000000-0000-0000-0000-000000000060") POLICY_ID = UUID("00000000-0000-0000-0000-000000000010") USER_ID = UUID("00000000-0000-0000-0000-000000000070") -CREATED_AT = datetime(2026, 1, 1, tzinfo=timezone.utc) +CREATED_AT = datetime(2026, 1, 1, tzinfo=UTC) def _association_read( @@ -122,7 +121,7 @@ def fallback(resource: str): readiness_probe=lambda: True, gateway_client_factory=lambda: None, metadata_reader_factory=lambda: None, - specification_provider=lambda: {}, + specification_provider=dict, v2_user_policy_service_factory=lambda: service, ) settings = RouteImplementationSettings( diff --git a/tests/unit/v2/test_user_policy_service.py b/tests/unit/v2/test_user_policy_service.py index 3444bf9ae..e7cb13dd0 100644 --- a/tests/unit/v2/test_user_policy_service.py +++ b/tests/unit/v2/test_user_policy_service.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime from uuid import UUID import pytest @@ -11,6 +11,7 @@ from sqlmodel import Session, create_engine, select from policyengine_api.data.v2.models import ( + V2_METADATA, LegacyUserPolicyMapping, Parameter, ParameterValue, @@ -19,7 +20,6 @@ TaxBenefitModelVersion, User, UserPolicy, - V2_METADATA, ) from policyengine_api.services.v2.user_policies.database_session import ( UserPolicyDatabaseSession, @@ -36,7 +36,6 @@ UserPolicyNotFoundError, ) - USER_ID = UUID("00000000-0000-0000-0000-000000000070") OTHER_USER_ID = UUID("00000000-0000-0000-0000-000000000071") @@ -76,7 +75,7 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: parameter=parameter, policy=policy, value_json=0.2, - start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), + start_date=datetime(2026, 1, 1, tzinfo=UTC), ) session.add_all( [