Skip to content

Commit

Permalink
test(providers/common/compat): add test case test_compat_has_only_ass…
Browse files Browse the repository at this point in the history
…et_methods and test_compat_has_asset_and_dataset_methods
  • Loading branch information
Lee-W committed Dec 5, 2024
1 parent 2b1c538 commit 71f3757
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 2 deletions.
4 changes: 2 additions & 2 deletions providers/src/airflow/providers/common/compat/lineage/hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,15 +90,15 @@ def get_hook_lineage_collector():
return _get_asset_compat_hook_lineage_collector()

# For the case that airflow has not yet upgraded to 2.10 or higher,
# but using the providers that alreay uses `get_hook_lineage_collector`
# but using the providers that already uses `get_hook_lineage_collector`
class NoOpCollector:
"""
NoOpCollector is a hook lineage collector that does nothing.
It is used when you want to disable lineage collection.
"""

# for providers that supports asset rename
# for providers that support asset rename
def add_input_asset(self, *_, **__):
pass

Expand Down
28 changes: 28 additions & 0 deletions providers/tests/common/compat/lineage/test_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 71f3757

Please sign in to comment.