from collections.abc import Sequence
from datetime import datetime, timedelta, timezone
from unittest import mock

import pytest
from dagster import (
    AssetCheckKey,
    AssetKey,
    AssetSpec,
    DagsterInstance,
    Definitions,
    SensorResult,
    asset,
    asset_check,
    build_sensor_context,
)
from dagster._core.definitions.assets.definition.assets_definition import AssetsDefinition
from dagster._core.definitions.events import AssetMaterialization
from dagster._core.definitions.materialize import materialize
from dagster._core.definitions.metadata.metadata_value import TimestampMetadataValue
from dagster._core.definitions.partitions.definition import DailyPartitionsDefinition
from dagster._core.definitions.sensor_definition import SensorEvaluationContext
from dagster._core.errors import DagsterInvariantViolationError
from dagster._core.execution.execute_in_process_result import ExecuteInProcessResult
from dagster._core.test_utils import freeze_time
from dagster._serdes import deserialize_value
from dagster._time import get_current_datetime
from dagster_airlift.constants import (
    DAG_ID_TAG_KEY,
    DAG_MAPPING_METADATA_KEY,
    DAG_RUN_ID_TAG_KEY,
    EFFECTIVE_TIMESTAMP_METADATA_KEY,
    TASK_ID_TAG_KEY,
)
from dagster_airlift.core import build_defs_from_airflow_instance
from dagster_airlift.core.airflow_defs_data import AirflowDefinitionsData
from dagster_airlift.core.sensor.sensor_builder import (
    AirflowPollingSensorCursor,
    AirliftSensorEventTransformerError,
)
from dagster_airlift.core.top_level_dag_def_api import assets_with_task_mappings
from dagster_airlift.test import make_dag_run, make_instance

from dagster_airlift_tests.unit_tests.conftest import (
    assert_expected_key_order,
    build_and_invoke_sensor,
    fully_loaded_repo_from_airflow_asset_graph,
)


def make_dag_key(dag_id: str) -> AssetKey:
    return AssetKey(["test_instance", "dag", dag_id])


def make_dag_key_str(dag_id: str) -> str:
    return make_dag_key(dag_id).to_user_string()


def test_dag_and_task_metadata(init_load_context: None, instance: DagsterInstance) -> None:
    """Test the metadata produced by a sensor for a single dag and task."""
    freeze_datetime = datetime(2021, 1, 1)

    with freeze_time(freeze_datetime):
        result, _ = build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", [])]},
            },
            instance=instance,
        )
        assert len(result.asset_events) == 2
        assert_expected_key_order(result.asset_events, ["a", make_dag_key_str("dag")])
        dag_mat = result.asset_events[1]
        expected_dag_metadata_keys = {
            "Airflow Run ID",
            "Start Date",
            "End Date",
            "Run Metadata (raw)",
            "Creation Timestamp",
            "Run Details",
            "Airflow Config",
            "Run Type",
            "dagster-airlift/effective-timestamp",
            "dagster-airlift/airflow-run-id",
        }
        assert set(dag_mat.metadata.keys()) == expected_dag_metadata_keys
        task_mat = result.asset_events[0]
        expected_task_metadata_keys = {
            "Airflow Run ID",
            "Start Date",
            "End Date",
            "Run Metadata (raw)",
            "Creation Timestamp",
            "Run Details",
            "Task Logs",
            "Airflow Config",
            "Run Type",
            "dagster-airlift/effective-timestamp",
            "dagster-airlift/airflow-run-id",
            "dagster-airlift/airflow-task-instance-logical-date",
        }
        assert set(task_mat.metadata.keys()) == expected_task_metadata_keys

        assert task_mat.metadata["Airflow Run ID"].value == "run-dag"
        assert (
            task_mat.metadata["Start Date"].value
            == (freeze_datetime - timedelta(minutes=10)).timestamp()
        )
        assert (
            task_mat.metadata["End Date"].value
            == (freeze_datetime - timedelta(seconds=1)).timestamp()
        )
        assert task_mat.metadata["Creation Timestamp"].value == freeze_datetime.timestamp()
        assert (
            task_mat.metadata["Run Details"].value
            == "[View Run](http://dummy.domain/dags/dag/grid?dag_run_id=run-dag&task_id=task)"
        )
        assert (
            task_mat.metadata["Task Logs"].value
            == "[View Logs](http://dummy.domain/dags/dag/grid?dag_run_id=run-dag&task_id=task&tab=logs)"
        )
        assert task_mat.metadata["Airflow Config"].value == {}
        assert task_mat.metadata["Run Type"].value == "manual"

        assert dag_mat.metadata["Airflow Run ID"].value == "run-dag"
        assert (
            dag_mat.metadata["Start Date"].value
            == (freeze_datetime - timedelta(minutes=10)).timestamp()
        )
        assert dag_mat.metadata["End Date"].value == freeze_datetime.timestamp()
        assert dag_mat.metadata["Creation Timestamp"].value == freeze_datetime.timestamp()
        assert (
            dag_mat.metadata["Run Details"].value
            == "[View Run](http://dummy.domain/dags/dag/grid?dag_run_id=run-dag&tab=details)"
        )
        assert dag_mat.metadata["Airflow Config"].value == {}
        assert dag_mat.metadata["Run Type"].value == "manual"


def test_interleaved_executions(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that the when task / dag completion is interleaved the correct ordering is preserved."""
    # Asset graph structure:
    #   a -> b where a and b are each in their own airflow tasks.
    #   c -> d where c and d are each in their own airflow tasks, in a different dag.
    freeze_datetime = datetime(2021, 1, 1)
    with freeze_time(freeze_datetime):
        result, context = build_and_invoke_sensor(
            assets_per_task={
                "dag1": {"task1": [("a", [])], "task2": [("b", ["a"])]},
                "dag2": {"task1": [("c", [])], "task2": [("d", ["c"])]},
            },
            instance=instance,
        )
        # We expect one asset materialization per asset.
        assert len(result.asset_events) == 6
        assert all(isinstance(event, AssetMaterialization) for event in result.asset_events)

        mats_order = [mat.asset_key.to_user_string() for mat in result.asset_events]
        # a should be before b
        assert mats_order.index("a") < mats_order.index("b")
        # c should be before d
        assert mats_order.index("c") < mats_order.index("d")
        # dag1 and dag2 should be after all task-mapped assets
        assert mats_order.index(make_dag_key_str("dag1")) >= 4
        assert mats_order.index(make_dag_key_str("dag2")) >= 4
        assert context.cursor
        cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert cursor.end_date_gte == freeze_datetime.timestamp()
        assert cursor.end_date_lte is None
        assert cursor.dag_query_offset == 0


def test_dependencies_within_tasks(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that a complex asset graph structure can be ingested in correct order from the sensor.
    Where a, b, and c are part of task 1, and d, e, and f are part of task 2.
    """
    # Asset graph structure:
    #   a
    #  / \
    # b   c
    #  \ /
    #   d
    #  / \
    # e   f
    freeze_datetime = datetime(2021, 1, 1)
    with freeze_time(freeze_datetime):
        result, context = build_and_invoke_sensor(
            assets_per_task={
                "dag": {
                    "task1": [("a", []), ("b", ["a"]), ("c", ["a"])],
                    "task2": [("d", ["b", "c"]), ("e", ["d"]), ("f", ["d"])],
                },
            },
            instance=instance,
        )
        assert len(result.asset_events) == 7
        assert_expected_key_order(
            result.asset_events, ["a", "b", "c", "d", "e", "f", make_dag_key_str("dag")]
        )
        assert context.cursor
        cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert cursor.end_date_gte == freeze_datetime.timestamp()
        assert cursor.end_date_lte is None
        assert cursor.dag_query_offset == 0


def test_outside_of_dag_dependency(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that if an asset has a transitive dependency on another asset within the same task, ordering is respected."""
    # a -> b -> c where a and c are in the same task, and b is not in any dag.
    freeze_datetime = datetime(2021, 1, 1)
    with freeze_time(freeze_datetime):
        result, context = build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", []), ("c", ["b"])]},
            },
            additional_defs=Definitions(assets=[AssetSpec(key="b", deps=["a"])]),
            instance=instance,
        )
        assert len(result.asset_events) == 3
        assert all(isinstance(event, AssetMaterialization) for event in result.asset_events)
        assert_expected_key_order(result.asset_events, ["a", "c", make_dag_key_str("dag")])
        assert context.cursor
        cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert cursor.end_date_gte == freeze_datetime.timestamp()
        assert cursor.end_date_lte is None
        assert cursor.dag_query_offset == 0


def test_request_asset_checks(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that when a new dag or task run is detected, a new check run is requested for all checks which may target that dag/task."""
    freeze_datetime = datetime(2021, 1, 1)

    dag_asset_key = make_dag_key("dag")

    @asset_check(asset="a")  # pyright: ignore[reportArgumentType]
    def check_task_asset():
        pass

    @asset_check(asset=dag_asset_key)  # pyright: ignore[reportArgumentType]
    def check_dag_asset():
        pass

    @asset_check(asset="c")  # pyright: ignore[reportArgumentType]
    def check_unrelated_asset():
        pass

    with freeze_time(freeze_datetime):
        result, context = build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", []), ("b", ["a"])]},
            },
            additional_defs=Definitions(
                asset_checks=[check_task_asset, check_dag_asset, check_unrelated_asset],
                assets=[AssetSpec(key="c")],
            ),
            instance=instance,
        )

        assert len(result.asset_events) == 3
        assert result.run_requests
        assert len(result.run_requests) == 1
        run_request = result.run_requests[0]
        assert run_request.asset_check_keys
        assert set(run_request.asset_check_keys) == {
            AssetCheckKey(name="check_task_asset", asset_key=AssetKey(["a"])),
            AssetCheckKey(name="check_dag_asset", asset_key=dag_asset_key),
        }
        assert context.cursor
        cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert cursor.end_date_gte == freeze_datetime.timestamp()
        assert cursor.end_date_lte is None
        assert cursor.dag_query_offset == 0


_CALLCOUNT = [0]


def _create_datetime_mocker(iter_times: list[datetime]):
    def _mock_get_current_datetime() -> datetime:
        the_time = iter_times[_CALLCOUNT[0]]
        _CALLCOUNT[0] += 1
        return the_time

    return _mock_get_current_datetime


def test_cursor(init_load_context: None, instance: DagsterInstance) -> None:
    """Test expected cursor behavior for sensor."""
    asset_and_dag_structure = {
        "dag1": {"task1": [("a", [])]},
        "dag2": {"task1": [("b", [])]},
    }

    with freeze_time(datetime(2021, 1, 1, tzinfo=timezone.utc)):
        # First, run through a full successful iteration of the sensor.
        # Expect time to move forward, and polled dag id to be None, since we completed iteration of all dags.
        # Then, run through a partial iteration of the sensor. We mock get_current_datetime to return a time
        # after timeout passes iteration start after the first call, meaning we should pause iteration.
        repo_def = fully_loaded_repo_from_airflow_asset_graph(asset_and_dag_structure)
        sensor = next(iter(repo_def.sensor_defs))
        context = build_sensor_context(repository_def=repo_def, instance=instance)
        result = sensor(context)
        assert isinstance(result, SensorResult)
        assert context.cursor
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert new_cursor.end_date_gte == datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.end_date_lte is None
        assert new_cursor.dag_query_offset == 0

    # Now, we expect that we will not have completed iteration before we need to pause evaluation.
    datetimes = [
        datetime(2021, 2, 1, tzinfo=timezone.utc),  # set initial time
        datetime(2021, 2, 1, 0, 0, 30, tzinfo=timezone.utc),  # initial iteration time
        datetime(
            2022, 2, 2, tzinfo=timezone.utc
        ),  # second iteration time, at which iteration should be paused
    ]
    with mock.patch(
        "dagster._time._mockable_get_current_datetime", wraps=_create_datetime_mocker(datetimes)
    ):
        result = sensor(context)
        assert isinstance(result, SensorResult)
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        # We didn't advance to the next effective timestamp, since we didn't complete iteration
        assert new_cursor.end_date_gte == datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
        # We have not yet moved forward
        assert new_cursor.end_date_lte == datetime(2021, 2, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.dag_query_offset == 1

        _CALLCOUNT[0] = 0
        # We weren't able to complete iteration, so we should pause iteration again
        result = sensor(context)
        assert isinstance(result, SensorResult)
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert new_cursor.end_date_gte == datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.end_date_lte == datetime(2021, 2, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.dag_query_offset == 2

        _CALLCOUNT[0] = 0
        # Now it should finish iteration.
        result = sensor(context)
        assert isinstance(result, SensorResult)
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert new_cursor.end_date_gte == datetime(2021, 2, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.end_date_lte is None
        assert new_cursor.dag_query_offset == 0


def test_legacy_cursor(init_load_context: None, instance: DagsterInstance) -> None:
    """Test the case where a legacy/uninterpretable cursor is provided to the sensor execution."""
    freeze_datetime = datetime(2021, 1, 1, tzinfo=timezone.utc)
    with freeze_time(freeze_datetime):
        repo_def = fully_loaded_repo_from_airflow_asset_graph({"dag": {"task": [("a", [])]}})
        sensor = next(iter(repo_def.sensor_defs))
        context = build_sensor_context(
            repository_def=repo_def, cursor=str(freeze_datetime.timestamp()), instance=instance
        )
        result = sensor(context)
        assert isinstance(result, SensorResult)
        assert context.cursor
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert new_cursor.end_date_gte == datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.end_date_lte is None
        assert new_cursor.dag_query_offset == 0


def test_no_runs(init_load_context: None, instance: DagsterInstance) -> None:
    """Test the case with no runs."""
    freeze_datetime = datetime(2021, 1, 1, tzinfo=timezone.utc)
    with freeze_time(freeze_datetime):
        repo_def = fully_loaded_repo_from_airflow_asset_graph(
            {"dag": {"task": [("a", [])]}},
            create_runs=False,
        )
        sensor = next(iter(repo_def.sensor_defs))
        context = build_sensor_context(repository_def=repo_def, instance=instance)
        result = sensor(context)
        assert isinstance(result, SensorResult)
        assert context.cursor
        new_cursor = deserialize_value(context.cursor, AirflowPollingSensorCursor)
        assert new_cursor.end_date_gte == datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
        assert new_cursor.end_date_lte is None
        assert new_cursor.dag_query_offset == 0
        assert not result.asset_events
        assert not result.run_requests


def simulate_materialize_from_proxy_operator(
    *,
    instance: DagsterInstance,
    assets_def: AssetsDefinition,
    dag_run_id: str,
    dag_id: str,
    task_id: str | None,
) -> ExecuteInProcessResult:
    # TODO consolidate with code in proxy operator
    tags = {
        DAG_ID_TAG_KEY: dag_id,
        DAG_RUN_ID_TAG_KEY: dag_run_id,
    }
    if task_id:
        tags[TASK_ID_TAG_KEY] = task_id
    return materialize(assets=[assets_def], instance=instance, tags=tags)


def test_pluggable_transformation(init_load_context: None, instance: DagsterInstance) -> None:
    """Test the case where a custom transformation is provided to the sensor."""

    def pluggable_event_transformer(
        context: SensorEvaluationContext,
        airflow_data: AirflowDefinitionsData,
        events: Sequence[AssetMaterialization],
    ) -> Sequence[AssetMaterialization]:
        assert isinstance(context, SensorEvaluationContext)
        assert isinstance(airflow_data, AirflowDefinitionsData)
        # Change the timestamp, which should also change the order. We expect this to be respected by the sensor.
        new_events = []
        for event in events:
            if AssetKey(["a"]) == event.asset_key:
                new_events.append(
                    event._replace(
                        metadata={
                            "test": "test",
                            EFFECTIVE_TIMESTAMP_METADATA_KEY: TimestampMetadataValue(1.0),
                        }
                    )
                )
            elif AssetKey.from_user_string(make_dag_key_str("dag")) == event.asset_key:
                new_events.append(
                    event._replace(
                        metadata={
                            "test": "test",
                            EFFECTIVE_TIMESTAMP_METADATA_KEY: TimestampMetadataValue(0.0),
                        }
                    )
                )
        return new_events

    result, _context = build_and_invoke_sensor(
        assets_per_task={
            "dag": {"task": [("a", [])]},
        },
        instance=instance,
        event_transformer_fn=pluggable_event_transformer,
    )
    assert len(result.asset_events) == 2
    assert_expected_key_order(result.asset_events, [make_dag_key_str("dag"), "a"])
    for event in result.asset_events:
        assert set(event.metadata.keys()) == {"test", EFFECTIVE_TIMESTAMP_METADATA_KEY}


def test_user_code_error_pluggable_transformation(
    init_load_context: None, instance: DagsterInstance
) -> None:
    """Test the case where a custom transformation is provided to the sensor, and the user code raises an error."""

    def pluggable_event_transformer(
        context: SensorEvaluationContext,
        airflow_data: AirflowDefinitionsData,
        events: Sequence[AssetMaterialization],
    ) -> Sequence[AssetMaterialization]:
        raise ValueError("User code error")

    with pytest.raises(AirliftSensorEventTransformerError):
        build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", [])]},
            },
            instance=instance,
            event_transformer_fn=pluggable_event_transformer,
        )


def test_missing_effective_timestamp_pluggable_impl(
    init_load_context: None, instance: DagsterInstance
) -> None:
    """Test the case where a custom transformation is provided to the sensor, and the user doesn't include effective timestamp metadata."""

    def missing_effective_timestamp(
        context: SensorEvaluationContext,
        airflow_data: AirflowDefinitionsData,
        events: Sequence[AssetMaterialization],
    ) -> Sequence[AssetMaterialization]:
        return [event._replace(metadata={}) for event in events]

    with pytest.raises(DagsterInvariantViolationError):
        build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", [])]},
            },
            instance=instance,
            event_transformer_fn=missing_effective_timestamp,
        )


def test_nonsense_result_pluggable_impl(init_load_context: None, instance: DagsterInstance) -> None:
    """Test the case where a nonsense result is returned from the custom transformation."""

    def nonsense_result(
        context: SensorEvaluationContext,
        airflow_data: AirflowDefinitionsData,
        events: Sequence[AssetMaterialization],
    ) -> Sequence[AssetMaterialization]:
        return [1, 2, 3]  # type: ignore # intentionally wrong

    with pytest.raises(DagsterInvariantViolationError):
        build_and_invoke_sensor(
            assets_per_task={
                "dag": {"task": [("a", [])]},
            },
            instance=instance,
            event_transformer_fn=nonsense_result,
        )


def test_dag_level_override_materializations(
    init_load_context: None, instance: DagsterInstance
) -> None:
    """Test that when a dag level override is provided, the sensor adds the expected asset materializations, in the correct order."""
    freeze_datetime = datetime(2021, 1, 1)

    with freeze_time(freeze_datetime):
        result, _context = build_and_invoke_sensor(
            assets_per_task={},
            dag_level_asset_overrides={"dag": ["a", "b"]},
            instance=instance,
        )
        assert len(result.asset_events) == 3  # 2 for the dag level overrides, 1 for the dag itself
        assert set([event.asset_key.to_user_string() for event in result.asset_events]) == {
            "a",
            "b",
            make_dag_key_str("dag"),
        }
        key_order = [event.asset_key.to_user_string() for event in result.asset_events]
        assert key_order.index("a") < key_order.index(make_dag_key_str("dag"))
        assert key_order.index("b") < key_order.index(make_dag_key_str("dag"))


def test_dag_level_override_existing_runs(
    init_load_context: None, instance: DagsterInstance
) -> None:
    """Test that when there is a run mapping to the same dag, the sensor doesn't add synthetic materializations to mapped assets (only the dag-level peered asset)."""
    freeze_datetime = datetime(2021, 1, 1)

    with freeze_time(freeze_datetime):
        dag_id = "dag"
        dag_run_id = f"run-{dag_id}"
        dag_runs = [
            make_dag_run(
                dag_id=dag_id,
                run_id=f"run-{dag_id}",
                start_date=get_current_datetime() - timedelta(minutes=10),
                end_date=get_current_datetime(),
            )
        ]

        airflow_instance = make_instance(
            dag_and_task_structure={dag_id: ["task"]},
            dag_runs=dag_runs,
            instance_name="test_instance",
        )

        @asset(metadata={DAG_MAPPING_METADATA_KEY: [{"dag_id": dag_id}]})
        def explicit_asset1() -> None:
            pass

        @asset(metadata={DAG_MAPPING_METADATA_KEY: [{"dag_id": dag_id}]})
        def explicit_asset2() -> None:
            pass

        # Simulate a proxied execution for the assets, meaning we don't need to synthesize materializations
        for assets_def in [explicit_asset1, explicit_asset2]:
            assert simulate_materialize_from_proxy_operator(
                instance=instance,
                assets_def=assets_def,
                dag_id=dag_id,
                task_id=None,
                dag_run_id=dag_run_id,
            ).success

        defs = build_defs_from_airflow_instance(
            airflow_instance=airflow_instance,
            defs=Definitions(
                assets=[explicit_asset1, explicit_asset2],
            ),
        )
        assert defs.sensors
        sensor = next(iter(defs.sensors))
        sensor_context = build_sensor_context(
            instance=instance, repository_def=defs.get_repository_def()
        )
        result = sensor(sensor_context)
        assert isinstance(result, SensorResult)
        assert len(result.asset_events) == 1
        assert result.asset_events[0].asset_key.to_user_string() == make_dag_key_str(dag_id)


def test_default_time_partitioned_asset(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that a task instance for a time-partitioned asset is correctly ingested."""
    defs = build_defs_from_airflow_instance(
        airflow_instance=make_instance(
            dag_and_task_structure={
                "dag": ["task"],
            },
            dag_runs=[
                make_dag_run(
                    dag_id="dag",
                    run_id="run-dag",
                    start_date=datetime(2021, 1, 2, 5, tzinfo=timezone.utc),
                    end_date=datetime(2021, 1, 2, 6, tzinfo=timezone.utc),
                    logical_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
                )
            ],
        ),
        defs=Definitions(
            assets=assets_with_task_mappings(
                dag_id="dag",
                task_mappings={
                    "task": [
                        AssetSpec(
                            key="a",
                            partitions_def=DailyPartitionsDefinition(
                                start_date=datetime(2021, 1, 1, tzinfo=timezone.utc)
                            ),
                        )
                    ],
                },
            ),
        ),
    )
    assert defs.sensors
    sensor = next(iter(defs.sensors))
    with freeze_time(datetime(2021, 1, 2, 6, tzinfo=timezone.utc)):
        sensor_context = build_sensor_context(
            repository_def=defs.get_repository_def(), instance=instance
        )
        result = sensor(sensor_context)
        assert isinstance(result, SensorResult)
        assert len(result.asset_events) == 2
        assert_expected_key_order(result.asset_events, ["a", "test_instance/dag/dag"])
        a_asset_mat = result.asset_events[0]
        assert isinstance(a_asset_mat, AssetMaterialization)
        # We expect the partition to match the logical date.
        assert a_asset_mat.partition == "2021-01-01"


def test_before_start_of_partitioned_asset(
    init_load_context: None, instance: DagsterInstance
) -> None:
    """We expect to throw an error if there is no matching partition after the start date."""
    defs = build_defs_from_airflow_instance(
        airflow_instance=make_instance(
            dag_and_task_structure={
                "dag": ["task"],
            },
            dag_runs=[
                make_dag_run(
                    dag_id="dag",
                    run_id="run-dag",
                    start_date=datetime(2021, 1, 1, 5, tzinfo=timezone.utc),
                    end_date=datetime(2021, 1, 1, 6, tzinfo=timezone.utc),
                    logical_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
                )
            ],
        ),
        defs=Definitions(
            assets=assets_with_task_mappings(
                dag_id="dag",
                task_mappings={
                    "task": [
                        AssetSpec(
                            key="a",
                            partitions_def=DailyPartitionsDefinition(
                                # Partitions definition starts after the logical date.
                                start_date=datetime(2021, 1, 2, tzinfo=timezone.utc)
                            ),
                        )
                    ],
                },
            ),
        ),
    )
    assert defs.sensors
    sensor = next(iter(defs.sensors))
    with freeze_time(datetime(2021, 1, 1, 6, tzinfo=timezone.utc)):
        sensor_context = build_sensor_context(
            repository_def=defs.get_repository_def(), instance=instance
        )
        with pytest.raises(AirliftSensorEventTransformerError):
            sensor(sensor_context)


def test_logical_date_mismatch(init_load_context: None, instance: DagsterInstance) -> None:
    """Test a logical date which does not align with the partition definition due to date mismatch."""
    defs = build_defs_from_airflow_instance(
        airflow_instance=make_instance(
            dag_and_task_structure={
                "dag": ["task"],
            },
            dag_runs=[
                make_dag_run(
                    dag_id="dag",
                    run_id="run-dag",
                    start_date=datetime(2021, 1, 1, 5, tzinfo=timezone.utc),
                    end_date=datetime(2021, 1, 1, 6, tzinfo=timezone.utc),
                    # Logical date isn't aligned with midnight.
                    logical_date=datetime(2021, 1, 1, 3, tzinfo=timezone.utc),
                )
            ],
        ),
        defs=Definitions(
            assets=assets_with_task_mappings(
                dag_id="dag",
                task_mappings={
                    "task": [
                        AssetSpec(
                            key="a",
                            partitions_def=DailyPartitionsDefinition(
                                start_date=datetime(2021, 1, 1, tzinfo=timezone.utc)
                            ),
                        )
                    ],
                },
            ),
        ),
    )
    assert defs.sensors
    sensor = next(iter(defs.sensors))
    with freeze_time(datetime(2021, 1, 1, 6, tzinfo=timezone.utc)):
        sensor_context = build_sensor_context(
            repository_def=defs.get_repository_def(), instance=instance
        )
        with pytest.raises(AirliftSensorEventTransformerError):
            sensor(sensor_context)


def test_partition_offset_mismatch(init_load_context: None, instance: DagsterInstance) -> None:
    """Test that partition offsets are respected when determining the partition corresponding to a logical date."""
    airflow_instance = make_instance(
        dag_and_task_structure={
            "dag": ["task"],
        },
        dag_runs=[
            make_dag_run(
                dag_id="dag",
                run_id="run-dag",
                start_date=datetime(2021, 1, 1, 5, tzinfo=timezone.utc),
                end_date=datetime(2021, 1, 1, 6, tzinfo=timezone.utc),
                # Logical date is 3 AM.
                logical_date=datetime(2021, 1, 1, 3, tzinfo=timezone.utc),
            )
        ],
    )

    defs = build_defs_from_airflow_instance(
        airflow_instance=airflow_instance,
        defs=Definitions(
            assets=assets_with_task_mappings(
                dag_id="dag",
                task_mappings={
                    "task": [
                        AssetSpec(
                            key="a",
                            partitions_def=DailyPartitionsDefinition(
                                start_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
                                hour_offset=6,
                            ),
                        )
                    ],
                },
            ),
        ),
    )
    # Due to differing hours offset, we expect an error to throw.
    assert defs.sensors
    sensor = next(iter(defs.sensors))
    with freeze_time(datetime(2021, 1, 1, 6, tzinfo=timezone.utc)):
        sensor_context = build_sensor_context(
            repository_def=defs.get_repository_def(), instance=instance
        )
        with pytest.raises(AirliftSensorEventTransformerError):
            sensor(sensor_context)

    # now, align the offset and expect success.
    defs = build_defs_from_airflow_instance(
        airflow_instance=airflow_instance,
        defs=Definitions(
            assets=assets_with_task_mappings(
                dag_id="dag",
                task_mappings={
                    "task": [
                        AssetSpec(
                            key="a",
                            partitions_def=DailyPartitionsDefinition(
                                start_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
                                hour_offset=3,
                            ),
                        )
                    ],
                },
            ),
        ),
    )
    assert defs.sensors
    sensor = next(iter(defs.sensors))
    with freeze_time(datetime(2021, 1, 1, 6, tzinfo=timezone.utc)):
        sensor_context = build_sensor_context(
            repository_def=defs.get_repository_def(), instance=instance
        )
        result = sensor(sensor_context)
        assert isinstance(result, SensorResult)
        assert len(result.asset_events) == 2
        assert_expected_key_order(result.asset_events, ["a", "test_instance/dag/dag"])
        a_asset_mat = result.asset_events[0]
        assert isinstance(a_asset_mat, AssetMaterialization)
        # We expect the partition to match the logical date.
        assert a_asset_mat.partition == "2021-01-01"
