diff --git a/providers/tests/common/compat/lineage/test_hook.py b/providers/tests/common/compat/lineage/test_hook.py index 1fd88405bc13f..13bb2cca7259c 100644 --- a/providers/tests/common/compat/lineage/test_hook.py +++ b/providers/tests/common/compat/lineage/test_hook.py @@ -16,9 +16,37 @@ # under the License. from __future__ import annotations +import pytest + from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector +from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS + def test_that_compat_does_not_raise(): # On compat tests this goes into ImportError code path assert get_hook_lineage_collector() is not None + assert get_hook_lineage_collector() is not None + + +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow 3.0+") +def test_compat_has_only_asset_methods(): + hook_lienage_collector = get_hook_lineage_collector() + + assert hook_lienage_collector.add_input_asset is not None + assert hook_lienage_collector.add_output_asset is not None + + with pytest.raises(AttributeError): + hook_lienage_collector.add_input_dataset + with pytest.raises(AttributeError): + hook_lienage_collector.add_output_dataset + + +@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow < 3.0") +def test_compat_has_asset_and_dataset_methods(): + hook_lienage_collector = get_hook_lineage_collector() + + assert hook_lienage_collector.add_input_asset is not None + assert hook_lienage_collector.add_output_asset is not None + assert hook_lienage_collector.add_input_dataset is not None + assert hook_lienage_collector.add_output_dataset is not None