diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f35f46c..2e34960 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -33,6 +33,10 @@ jobs: - run: make init + - name: Unit Tests + run: | + make test-unit + - name: Integration Tests run: | make test-integration diff --git a/Makefile b/Makefile index 8544844..e6481a5 100644 --- a/Makefile +++ b/Makefile @@ -27,7 +27,10 @@ compose-down:: docker compose down --rmi 'all' -test:: test-integration test-acceptance +test:: test-unit test-integration test-acceptance + +test-unit:: + python -m pytest tests/unit test-integration:: python -m pytest tests/integration diff --git a/dags/collection_config.py b/dags/collection_config.py index 4be4bfc..03d498e 100644 --- a/dags/collection_config.py +++ b/dags/collection_config.py @@ -7,10 +7,18 @@ defined here rather than being hardcoded into the DAG generation logic. """ +from datetime import datetime, timedelta, timezone from typing import Optional +from dateutil.rrule import rrulestr from pydantic import BaseModel +# fixed anchor for expanding schedule_rrule recurrences from; only matters for cadences whose +# phase depends on a start date (e.g. fortnightly) - harmless for DAILY/MONTHLY;BYDAY rules. +# Must be timezone-aware since Airflow's logical_date is UTC-aware, and dateutil can't compare +# naive and aware datetimes. +RRULE_SERIES_START = datetime(2025, 1, 1, tzinfo=timezone.utc) + class CollectionDagConfig(BaseModel): cpu: int = 8192 @@ -27,6 +35,9 @@ class CollectionDagConfig(BaseModel): # how often (in seconds) the ECS task operators poll CloudWatch for new log lines; too low a # value puts workers under load when many ECS tasks are triggered at once awslogs_fetch_interval_seconds: int = 10 + # RFC 5545 recurrence rule (RRULE) controlling which of the scheduler's daily runs should + # actually trigger this collection; "FREQ=DAILY" (the default) means every run + schedule_rrule: str = "FREQ=DAILY" DEFAULT_COLLECTION_CONFIG = CollectionDagConfig() @@ -34,11 +45,21 @@ class CollectionDagConfig(BaseModel): # collections whose defaults diverge from DEFAULT_COLLECTION_CONFIG COLLECTION_CONFIG_OVERRIDES = { # title-boundary's EMR job can otherwise consume all of the vCPU available to the - # shared EMR Serverless application, starving other collections' jobs - "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50), + # shared EMR Serverless application, starving other collections' jobs. It's also only + # scheduled for the first Monday of the month, a day after HM Land Registry's own release + # schedule for the INSPIRE Index Polygons data it's built from (first Sunday), so their + # data is available by the time we run + "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50, schedule_rrule="FREQ=MONTHLY;BYDAY=1MO"), } def get_collection_dag_config(collection: str) -> CollectionDagConfig: """Return the default param config for a collection, falling back to the generic defaults.""" return COLLECTION_CONFIG_OVERRIDES.get(collection, DEFAULT_COLLECTION_CONFIG) + + +def collection_schedule_matches(collection: str, logical_date, **_) -> bool: + """Whether a collection's schedule_rrule has an occurrence on logical_date's date.""" + rrule_str = get_collection_dag_config(collection).schedule_rrule + occurrence = rrulestr(rrule_str, dtstart=RRULE_SERIES_START).after(logical_date - timedelta(seconds=1), inc=True) + return occurrence is not None and occurrence.date() == logical_date.date() diff --git a/dags/dag_triggers.py b/dags/dag_triggers.py index 650920f..265e7b4 100644 --- a/dags/dag_triggers.py +++ b/dags/dag_triggers.py @@ -7,11 +7,24 @@ from datetime import datetime from airflow import DAG +from airflow.operators.python import ShortCircuitOperator from airflow.operators.trigger_dagrun import TriggerDagRunOperator from airflow.utils.trigger_rule import TriggerRule +from collection_config import DEFAULT_COLLECTION_CONFIG, collection_schedule_matches, get_collection_dag_config from collection_schema import CollectionSelection from utils import filter_collections_for_env, get_collections_dict, get_config, load_specification_datasets, sort_collections_dict +# title-boundary runs via its newer, EMR-based pipeline; every other collection still runs via +# the original collection DAG +NEW_COLLECTION_DAG_COLLECTIONS = {"title-boundary"} + + +def trigger_dag_id_for(collection: str) -> str: + if collection in NEW_COLLECTION_DAG_COLLECTIONS: + return f"new-{collection}-collection" + return f"{collection}-collection" + + config = get_config() dag_schedule = config.get("schedule", None) # Use "None" as a fallback if "schedule" key is missing dag_max_active_tasks = config.get("max_active_tasks") @@ -61,15 +74,31 @@ def collection_selected(collection_name, configuration): collection_dag = TriggerDagRunOperator( task_id=f"trigger-{collection}-collection-dag", - trigger_dag_id=f"{collection}-collection", + trigger_dag_id=trigger_dag_id_for(collection), wait_for_completion=True, trigger_rule=TriggerRule.ALL_DONE, priority_weight=CUSTOM_COLLECTION_DAG_WEIGHTING.get(collection, DEFAULT_WEIGHTING), conf=conf, ) + + # the scheduler runs every day; collections with a non-default schedule_rrule + # (e.g. title-boundary, monthly) only actually get triggered on a matching day + entry_task = collection_dag + collection_dag_config = get_collection_dag_config(collection) + if collection_dag_config.schedule_rrule != DEFAULT_COLLECTION_CONFIG.schedule_rrule: + check_schedule = ShortCircuitOperator( + task_id=f"check-{collection}-schedule", + python_callable=collection_schedule_matches, + op_kwargs={"collection": collection}, + trigger_rule=TriggerRule.ALL_DONE, + ignore_downstream_trigger_rules=False, + ) + check_schedule >> collection_dag + entry_task = check_schedule + collection_tasks.append(collection_dag) if organisation_collection_selected: - run_org_builder_dag >> collection_dag + run_org_builder_dag >> entry_task dlb_dag = TriggerDagRunOperator( task_id="trigger-digital-land-builder-dag", trigger_dag_id="build-digital-land-builder", wait_for_completion=True, trigger_rule=TriggerRule.ALL_DONE @@ -101,7 +130,9 @@ def collection_selected(collection_name, configuration): # Set custom CPU for listed-building collection conf = {"cpu": 16384, "transformed-jobs": 16} if collection in ("listed-building", "tree-preservation-order") else {} - collection_dag = TriggerDagRunOperator(task_id=f"trigger-{collection}-collection-dag", trigger_dag_id=f"{collection}-collection", wait_for_completion=True, conf=conf) + collection_dag = TriggerDagRunOperator( + task_id=f"trigger-{collection}-collection-dag", trigger_dag_id=trigger_dag_id_for(collection), wait_for_completion=True, conf=conf + ) collection_tasks.append(collection_dag) run_org_builder_dag >> collection_dag diff --git a/requirements/requirements.in b/requirements/requirements.in index 7dd9385..6e8c0b5 100644 --- a/requirements/requirements.in +++ b/requirements/requirements.in @@ -3,5 +3,6 @@ apache-airflow-providers-amazon apache-airflow-providers-slack boto pydantic +python-dateutil click utils \ No newline at end of file diff --git a/requirements/requirements.txt b/requirements/requirements.txt index c984697..fe2454e 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -385,6 +385,7 @@ python-daemon==3.0.1 # via apache-airflow python-dateutil==2.9.0.post0 # via + # -r requirements/requirements.in # apache-airflow # botocore # croniter diff --git a/tests/acceptance/test_dags_render.py b/tests/acceptance/test_dags_render.py index c7abe94..56c0fd1 100644 --- a/tests/acceptance/test_dags_render.py +++ b/tests/acceptance/test_dags_render.py @@ -9,6 +9,12 @@ def test_dag_rendering(): dag_bag = DagBag(dag_folder=DAG_FOLDER, include_examples=False) assert len(dag_bag.dags) > 0, "No DAGs found in DAG bag!" + # DagBag swallows per-file import errors rather than raising, so a broken DAG module + # wouldn't otherwise fail this test - it would just silently be missing from dag_bag.dags + if dag_bag.import_errors: + errors = "\n".join(f"{path}: {error}" for path, error in dag_bag.import_errors.items()) + pytest.fail(f"DAG files failed to import:\n{errors}") + # Collect any DAGs that failed to load failed_dags = [] diff --git a/tests/conftest.py b/tests/conftest.py index 0de7191..444fbac 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,44 +6,45 @@ import pytest from moto import mock_aws +CONFIG_PATH = Path("dags/config.json") +CONFIG_BACKUP_PATH = CONFIG_PATH.with_suffix(".json.backup") -@pytest.fixture(scope="session", autouse=True) -def create_test_config(): - """ - Create a test config.json file for DAG tests. +# Matches the output of: bin/generate_dag_config.py --env development +TEST_CONFIG = {"env": "development", "schedule": "0 0 * * *", "max_active_tasks": 50, "collection_selection": "explicit", "collections": ["ancient-woodland", "organisation"]} + +# Several DAG modules call get_secrets("emr_execution_role", env) at import time, which hits real +# AWS Secrets Manager unless overridden. aws_secrets_manager.get_secret_emr_compatible supports a +# SECRET_ environment variable escape hatch specifically for this - see its docstring. +SECRET_ENV_VAR = "SECRET__DEVELOPMENT_PD_BATCH_DEPLOYMENT_VARIABLES_SECRET" +TEST_SECRET_VALUE = json.dumps({"emr_execution_role": "arn:aws:iam::000000000000:role/test-emr-execution-role"}) - This fixture runs automatically before any tests and creates the config file - that DAGs expect to find in dags/config.json. If a config already exists, - it will be backed up and restored after tests complete. - The test configuration matches the output of: - bin/generate_dag_config.py --env development +def pytest_configure(config): + """ + Create a test config.json file, and fake the AWS secret DAGs read at import time. - To customize the test config, modify the test_config dictionary below. + Some tests import DAG modules directly at module scope (e.g. `from dags.digital_land_builder + import dag`), which happens during collection - before any pytest fixture would run. This + hook runs earlier still, so both are guaranteed to be in place before anything is collected. + If a config already exists, it's backed up and restored in pytest_unconfigure. """ - config_path = Path("dags/config.json") + if CONFIG_PATH.exists(): + CONFIG_PATH.rename(CONFIG_BACKUP_PATH) - # Create test configuration (matches development environment) - test_config = {"env": "development", "schedule": "0 0 * * *", "max_active_tasks": 50, "collection_selection": "explicit", "collections": ["ancient-woodland", "organisation"]} + with open(CONFIG_PATH, "w") as f: + json.dump(TEST_CONFIG, f, indent=4) - # Check if config already exists (backup if it does) - backup_path = None - if config_path.exists(): - backup_path = config_path.with_suffix(".json.backup") - config_path.rename(backup_path) + os.environ[SECRET_ENV_VAR] = TEST_SECRET_VALUE - # Write test config - with open(config_path, "w") as f: - json.dump(test_config, f, indent=4) - yield config_path +def pytest_unconfigure(config): + if CONFIG_PATH.exists(): + CONFIG_PATH.unlink() - # Cleanup: remove test config and restore backup if it existed - if config_path.exists(): - config_path.unlink() + if CONFIG_BACKUP_PATH.exists(): + CONFIG_BACKUP_PATH.rename(CONFIG_PATH) - if backup_path and backup_path.exists(): - backup_path.rename(config_path) + os.environ.pop(SECRET_ENV_VAR, None) @pytest.fixture(scope="function") diff --git a/tests/unit/test_collection_config.py b/tests/unit/test_collection_config.py new file mode 100644 index 0000000..8fc4e78 --- /dev/null +++ b/tests/unit/test_collection_config.py @@ -0,0 +1,51 @@ +import pendulum +import pytest + +from dags.collection_config import DEFAULT_COLLECTION_CONFIG, collection_schedule_matches, get_collection_dag_config + + +def test_get_collection_dag_config_returns_default_for_unknown_collection(): + assert get_collection_dag_config("central-activities-zone") is DEFAULT_COLLECTION_CONFIG + + +def test_get_collection_dag_config_returns_override_for_title_boundary(): + config = get_collection_dag_config("title-boundary") + assert config.schedule_rrule == "FREQ=MONTHLY;BYDAY=1MO" + assert config.transform_batch_size == 100 + assert config.max_executors == 50 + + +def test_default_collection_schedule_is_daily(): + assert DEFAULT_COLLECTION_CONFIG.schedule_rrule == "FREQ=DAILY" + + +def test_collection_schedule_matches_accepts_timezone_aware_logical_date(): + """Regression test: Airflow's logical_date is UTC-aware. dateutil raises "can't compare + offset-naive and offset-aware datetimes" if the rrule's own anchor isn't aware too.""" + logical_date = pendulum.datetime(2026, 8, 3, tz="UTC") + collection_schedule_matches("title-boundary", logical_date) # should not raise + + +@pytest.mark.parametrize( + "logical_date,expected", + [ + (pendulum.datetime(2026, 8, 3, tz="UTC"), True), # first Monday of August 2026 + (pendulum.datetime(2026, 8, 2, tz="UTC"), False), # first Sunday - not a match + (pendulum.datetime(2026, 8, 10, tz="UTC"), False), # second Monday - not a match + (pendulum.datetime(2026, 9, 7, tz="UTC"), True), # first Monday of September 2026 + ], +) +def test_collection_schedule_matches_title_boundary_first_monday(logical_date, expected): + assert collection_schedule_matches("title-boundary", logical_date) is expected + + +@pytest.mark.parametrize( + "logical_date", + [ + pendulum.datetime(2026, 8, 1, tz="UTC"), + pendulum.datetime(2026, 8, 15, tz="UTC"), + pendulum.datetime(2026, 12, 25, tz="UTC"), + ], +) +def test_collection_schedule_matches_default_daily_collection_always_matches(logical_date): + assert collection_schedule_matches("central-activities-zone", logical_date) is True