SameerMesiah97 commented on code in PR #77068:
URL: https://github.com/apache/airflow/pull/77068#discussion_r4569404082
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -2252,14 +2252,58 @@ def
test_get_conn_passes_workload_identity_provider_to_connect(self):
assert call_kwargs["authenticator"] == "WORKLOAD_IDENTITY"
assert call_kwargs["workload_identity_provider"] == "GCP"
- def
test_get_connection_form_widgets_exposes_workload_identity_provider(self):
- """The connection form must expose a workload_identity_provider field
so users can set it in the UI."""
+ @pytest.mark.parametrize("field", ["workload_identity_provider",
"authenticator"])
+ def test_get_connection_form_widgets_exposes_auth_fields(self, field):
+ """The connection form must expose auth fields so users can set them
in the UI."""
pytest.importorskip("flask_appbuilder")
pytest.importorskip("flask_babel")
widgets = SnowflakeHook.get_connection_form_widgets()
- assert "workload_identity_provider" in widgets
+ assert field in widgets
+
+ @pytest.mark.parametrize(
+ ("extra_authenticator", "hook_authenticator", "expected"),
+ [
+ pytest.param(None, None, "WORKLOAD_IDENTITY", id="inferred"),
+ pytest.param("", None, "WORKLOAD_IDENTITY",
id="inferred-from-blank-form-field"),
+ pytest.param("oauth", None, "oauth", id="explicit-extra-wins"),
+ pytest.param(None, "snowflake", "snowflake",
id="explicit-hook-kwarg-wins"),
Review Comment:
I would add the below test case to cover the override behaviour:
```
pytest.param(
"oauth", "snowflake", "snowflake",
id="hook-overrides-extra",
),
```
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -2252,14 +2252,58 @@ def
test_get_conn_passes_workload_identity_provider_to_connect(self):
assert call_kwargs["authenticator"] == "WORKLOAD_IDENTITY"
assert call_kwargs["workload_identity_provider"] == "GCP"
- def
test_get_connection_form_widgets_exposes_workload_identity_provider(self):
- """The connection form must expose a workload_identity_provider field
so users can set it in the UI."""
+ @pytest.mark.parametrize("field", ["workload_identity_provider",
"authenticator"])
+ def test_get_connection_form_widgets_exposes_auth_fields(self, field):
+ """The connection form must expose auth fields so users can set them
in the UI."""
pytest.importorskip("flask_appbuilder")
pytest.importorskip("flask_babel")
widgets = SnowflakeHook.get_connection_form_widgets()
- assert "workload_identity_provider" in widgets
+ assert field in widgets
+
+ @pytest.mark.parametrize(
+ ("extra_authenticator", "hook_authenticator", "expected"),
+ [
+ pytest.param(None, None, "WORKLOAD_IDENTITY", id="inferred"),
+ pytest.param("", None, "WORKLOAD_IDENTITY",
id="inferred-from-blank-form-field"),
+ pytest.param("oauth", None, "oauth", id="explicit-extra-wins"),
+ pytest.param(None, "snowflake", "snowflake",
id="explicit-hook-kwarg-wins"),
+ ],
+ )
+ def test_get_conn_params_infers_workload_identity_authenticator(
+ self, extra_authenticator, hook_authenticator, expected
+ ):
+ connection_kwargs = deepcopy(BASE_CONNECTION_KWARGS)
+ connection_kwargs["extra"]["workload_identity_provider"] = "GCP"
+ if extra_authenticator is not None:
+ connection_kwargs["extra"]["authenticator"] = extra_authenticator
+
+ with mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()):
+ hook = SnowflakeHook(snowflake_conn_id="test_conn",
authenticator=hook_authenticator)
+ conn_params = hook._get_static_conn_params
+
+ assert conn_params["authenticator"] == expected
+
+ @pytest.mark.parametrize(
+ ("extra", "expected"),
+ [
+ pytest.param({"authenticator": ""}, "snowflake",
id="blank-form-field-uses-default"),
+ pytest.param(
+ {"extra__snowflake__authenticator":
"programmatic_access_token"},
+ "programmatic_access_token",
+ id="backcompat-prefix",
+ ),
Review Comment:
You coud add a case with both `workload_identity_provider="GCP"` and
`extra__snowflake__authenticator="programmatic_access_token"`, asserting that
the authenticator remains `"programmatic_access_token"`. The existing
prefixed-key test doesn’t configure WIF, so it doesn’t verify that an explicit
prefixed authenticator prevents WIF inference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
