diff --git a/cosmos/profiles/databricks/oauth.py b/cosmos/profiles/databricks/oauth.py index fd6875c497..a2ef66d4e0 100644 --- a/cosmos/profiles/databricks/oauth.py +++ b/cosmos/profiles/databricks/oauth.py @@ -36,6 +36,14 @@ class DatabricksOauthProfileMapping(BaseProfileMapping): "http_path": "extra.http_path", } + @property + def mock_profile(self) -> dict[str, Any | None]: + """Generates a mock profile.""" + return { + **super().mock_profile, + "auth_type": "oauth", + } + @property def profile(self) -> dict[str, Any | None]: """Generates profile. The client-id and client-secret is stored in an environment variable.""" diff --git a/tests/profiles/databricks/test_dbr_oauth.py b/tests/profiles/databricks/test_dbr_oauth.py index fe6029508f..2df653e089 100644 --- a/tests/profiles/databricks/test_dbr_oauth.py +++ b/tests/profiles/databricks/test_dbr_oauth.py @@ -69,3 +69,11 @@ def test_connection_claiming() -> None: with patch("cosmos.profiles.base.BaseHook.get_connection", return_value=conn): profile_mapping = DatabricksOauthProfileMapping(conn, {"schema": "my_schema"}) assert profile_mapping.can_claim_connection() + + +def test_mock_profile() -> None: + """ + Tests that the Databricks mock profile is generated correctly. + """ + profile_mapping = DatabricksOauthProfileMapping("conn_id", {"schema": "my_schema"}) + assert profile_mapping.mock_profile.get("auth_type") == "oauth"