o-nikolas commented on code in PR #55694:
URL: https://github.com/apache/airflow/pull/55694#discussion_r2395489120
##########
providers/amazon/tests/unit/amazon/aws/hooks/test_athena_sql.py:
##########
@@ -153,3 +163,123 @@ def test_conn_id_override_setter(self):
hook = AthenaSQLHook(athena_conn_id=AWS_ATHENA_CONN_ID,
aws_conn_id=AWS_CONN_ID)
assert hook.athena_conn_id == AWS_ATHENA_CONN_ID
assert hook.aws_conn_id == AWS_CONN_ID
+
+ def test_hook_params_handling(self):
+ """Test that hook_params are properly handled and don't cause
TypeError."""
+ # Test that hook_params with Athena-specific parameters don't cause
errors
+ hook = AthenaSQLHook(
+ athena_conn_id="test_conn",
+ s3_staging_dir="s3://test-bucket/staging/",
+ work_group="test-workgroup",
+ driver="rest",
+ aws_domain="amazonaws.com",
+ session_kwargs={"profile_name": "test"},
+ config_kwargs={"retries": {"max_attempts": 5}},
+ role_arn="arn:aws:iam::123456789012:role/test-role",
+ assume_role_method="assume_role",
+ assume_role_kwargs={"RoleSessionName": "airflow-test"},
+ aws_session_token="test-token",
+ endpoint_url="https://athena.us-east-1.amazonaws.com";,
+ )
+
+ # Verify that the parameters were extracted correctly
+ assert hook.s3_staging_dir == "s3://test-bucket/staging/"
+ assert hook.work_group == "test-workgroup"
+ assert hook.driver == "rest"
+ assert hook.aws_domain == "amazonaws.com"
+ assert hook.session_kwargs == {"profile_name": "test"}
+ assert hook.config_kwargs == {"retries": {"max_attempts": 5}}
+ assert hook.role_arn == "arn:aws:iam::123456789012:role/test-role"
+ assert hook.assume_role_method == "assume_role"
+ assert hook.assume_role_kwargs == {"RoleSessionName": "airflow-test"}
+ assert hook.aws_session_token == "test-token"
+ assert hook.endpoint_url == "https://athena.us-east-1.amazonaws.com";
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.athena_sql.pyathena.connect")
+
@mock.patch("airflow.providers.amazon.aws.hooks.athena_sql.AthenaSQLHook.get_session")
+ def test_get_conn_with_hook_params(self, mock_get_session, mock_connect):
+ """Test that get_conn uses hook_params when provided."""
+ # Create hook with hook_params
+ hook = AthenaSQLHook(
+ athena_conn_id="test_conn",
+ s3_staging_dir="s3://test-bucket/staging/",
+ work_group="test-workgroup",
+ )
+
+ # Mock the connection
+ conn = Connection(
+ conn_type="athena",
+ schema="test_schema",
+ extra={"region_name": "us-east-1"},
+ )
+ hook.get_connection = mock.Mock(return_value=conn)
+
+ # Call get_conn
+ hook.get_conn()
Review Comment:
I still think some that are left are not helpful. But I'm not going to burn
more time on this PR.
--
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]
