shahar1 commented on code in PR #74804:
URL: https://github.com/apache/airflow/pull/74804#discussion_r4302545605
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -1567,6 +1568,118 @@ def test_get_oauth_token_with_scope(self,
mock_requests_post):
assert called_data["scope"] == "default"
assert called_data["grant_type"] == "client_credentials"
+ @mock.patch("requests.post")
+ def test_get_oauth_token_with_password_grant_type(self,
mock_requests_post):
+ """
+ Verify that `get_oauth_token` returns an access token and includes the
+ provided required parameters for password type of grant in the
outgoing OAuth request payload.
+ """
+
+ mock_requests_post.return_value = Mock(
+ spec=requests.Response,
+ status_code=200,
+ json=lambda: {"access_token": "dummy_token", "expires_in": 600},
+ )
+
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "login": "client_id",
+ "password": "client_secret",
+ "extra": {
+ "account": "airflow",
+ "authenticator": "oauth",
+ "grant_type": "password",
+ "token_endpoint": "https://localhost";,
+ "oauth_username": "user1",
+ "oauth_password": "password1",
+ },
+ }
+
+ with mock.patch.dict(
+ "os.environ",
+ {"AIRFLOW_CONN_TEST_CONN":
Connection(**connection_kwargs).get_uri()},
+ ):
+ hook = SnowflakeHook(snowflake_conn_id="test_conn")
+ token = hook.get_oauth_token(grant_type="password")
+
+ assert token == "dummy_token"
+
+ called_data = mock_requests_post.call_args.kwargs["data"]
+
+ assert called_data["grant_type"] == "password"
+ assert called_data["username"] == "user1"
+ assert called_data["password"] == "password1"
+
+ def
test_get_conn_params_should_fail_on_lack_of_oauth_password_for_grant_type_password(self):
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "password": None,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "authenticator": "oauth",
+ "grant_type": "password",
+ "token_endpoint": "https://localhost";,
+ "oauth_username": "test1",
+ },
+ }
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ pytest.raises(
+ ValueError, match="Resource owner user password must be
provided for password grant type"
+ ),
+ ):
+ SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
+ def
test_get_conn_params_should_fail_on_lack_of_oauth_username_for_grant_type_password(self):
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "password": None,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "authenticator": "oauth",
+ "token_endpoint": "https://localhost";,
+ "grant_type": "password",
+ "oauth_password": "pass1",
+ },
+ }
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ pytest.raises(
+ ValueError, match="Resource owner user name must be provided
for password grant type"
+ ),
+ ):
+ SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
+ def
test_get_conn_params_should_fail_on_lack_of_token_endpoint_for_grant_type_password(self):
Review Comment:
Could you please consolidate these test into a single parametrized test?
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -1567,6 +1568,118 @@ def test_get_oauth_token_with_scope(self,
mock_requests_post):
assert called_data["scope"] == "default"
assert called_data["grant_type"] == "client_credentials"
+ @mock.patch("requests.post")
+ def test_get_oauth_token_with_password_grant_type(self,
mock_requests_post):
+ """
+ Verify that `get_oauth_token` returns an access token and includes the
+ provided required parameters for password type of grant in the
outgoing OAuth request payload.
+ """
+
+ mock_requests_post.return_value = Mock(
+ spec=requests.Response,
+ status_code=200,
+ json=lambda: {"access_token": "dummy_token", "expires_in": 600},
+ )
+
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "login": "client_id",
+ "password": "client_secret",
+ "extra": {
+ "account": "airflow",
+ "authenticator": "oauth",
+ "grant_type": "password",
+ "token_endpoint": "https://localhost";,
+ "oauth_username": "user1",
+ "oauth_password": "password1",
+ },
+ }
+
+ with mock.patch.dict(
+ "os.environ",
+ {"AIRFLOW_CONN_TEST_CONN":
Connection(**connection_kwargs).get_uri()},
+ ):
+ hook = SnowflakeHook(snowflake_conn_id="test_conn")
+ token = hook.get_oauth_token(grant_type="password")
+
+ assert token == "dummy_token"
+
+ called_data = mock_requests_post.call_args.kwargs["data"]
+
+ assert called_data["grant_type"] == "password"
+ assert called_data["username"] == "user1"
+ assert called_data["password"] == "password1"
+
+ def
test_get_conn_params_should_fail_on_lack_of_oauth_password_for_grant_type_password(self):
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "password": None,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "authenticator": "oauth",
+ "grant_type": "password",
+ "token_endpoint": "https://localhost";,
+ "oauth_username": "test1",
+ },
+ }
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ pytest.raises(
+ ValueError, match="Resource owner user password must be
provided for password grant type"
+ ),
+ ):
+ SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
+ def
test_get_conn_params_should_fail_on_lack_of_oauth_username_for_grant_type_password(self):
+ connection_kwargs = {
+ **BASE_CONNECTION_KWARGS,
+ "password": None,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "authenticator": "oauth",
+ "token_endpoint": "https://localhost";,
+ "grant_type": "password",
+ "oauth_password": "pass1",
+ },
+ }
+ with (
+ mock.patch.dict("os.environ",
AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()),
+ pytest.raises(
+ ValueError, match="Resource owner user name must be provided
for password grant type"
+ ),
+ ):
+ SnowflakeHook(snowflake_conn_id="test_conn").get_conn()
+
+ def
test_get_conn_params_should_fail_on_lack_of_token_endpoint_for_grant_type_password(self):
Review Comment:
Could you please consolidate these tests into a single parametrized test?
--
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]
