diff --git a/cosmos/plugin/__init__.py b/cosmos/plugin/__init__.py index 84663663fb..d91f5cb6f7 100644 --- a/cosmos/plugin/__init__.py +++ b/cosmos/plugin/__init__.py @@ -1,11 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + from airflow import __version__ as airflow_version from packaging import version from cosmos.constants import _AIRFLOW3_MAJOR_VERSION +if TYPE_CHECKING: # pragma: no cover + from .plugin_impl import CosmosPlugin as _CosmosPluginType + +CosmosPlugin: _CosmosPluginType | None = None # The plugin is only loaded if the Airflow version is less than 3.0. This is because the plugin is incompatible with # Airflow 3.0 and above. Once the compatibility issue is resolved as part of # https://github.com/astronomer/astronomer-cosmos/issues/1587, the import statement can be moved outside of the # conditional block. if version.parse(airflow_version).major < _AIRFLOW3_MAJOR_VERSION: - from .plugin_impl import CosmosPlugin as CosmosPlugin + from .plugin_impl import CosmosPlugin as CosmosPlugin # type: ignore[assignment] diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index 93b7d4b7c7..9dd385a77e 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -380,3 +380,7 @@ def test_has_access_with_permissions_in_astro_must_include_custom_menu(url_path, mock_check_auth = cosmos.plugin.plugin_impl.dbt_docs_view.appbuilder.sm.check_authorization app.get(url_path) assert mock_check_auth.call_args[0][0] == [("menu_access", "Custom Menu"), ("can_read", "Website")] + + +def test_cosmos_plugin_enabled_on_airflow2(): + assert cosmos.plugin.CosmosPlugin is not None diff --git a/tests/plugin/test_plugin_af3.py b/tests/plugin/test_plugin_af3.py new file mode 100644 index 0000000000..929fcbfc7d --- /dev/null +++ b/tests/plugin/test_plugin_af3.py @@ -0,0 +1,16 @@ +# This module is added specifically to test that the plugin is disabled in Airflow 3.x. +# Currently, we skip the entire test_plugin.py module since it contains multiple tests that are only compatible with Airflow 2.x. +# Adding this test to that file would require selectively applying pytest.mark.skipif to each individual test, +# which would add unnecessary complexity. +# To keep things straightforward, we isolate the one relevant test for Airflow 3 here, +# while continuing to skip the existing plugin test module entirely for Airflow 3.x. +import pytest +from airflow import __version__ as airflow_version +from packaging import version + +import cosmos.plugin + + +@pytest.mark.skipif(version.parse(airflow_version).major == 2, reason="Plugin is available in Airflow 2.x") +def test_cosmos_plugin_disabled_on_airflow3(): + assert cosmos.plugin.CosmosPlugin is None