Lee-W commented on code in PR #56324:
URL: https://github.com/apache/airflow/pull/56324#discussion_r3091879237
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -275,11 +301,26 @@ def __init__(
self,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
+ quota_project_id: str | None = None,
Review Comment:
```suggestion
impersonation_chain: str | Sequence[str] | None = None,
*,
quota_project_id: str | None = None,
```
nit: the order does not have any special meaning here. thus, keyword only
arg would be better.
##########
providers/google/tests/unit/google/common/hooks/test_base_google.py:
##########
@@ -541,6 +541,59 @@ def
test_get_credentials_and_project_id_with_default_auth_and_overridden_project
)
assert result == ("CREDENTIALS", "SECOND_PROJECT_ID")
+ def test_quota_project_id_init(self):
+ """Test that quota project ID is properly initialized."""
+ hook = GoogleBaseHook(gcp_conn_id="google_cloud_default",
quota_project_id="test-quota-project")
+ assert hook.quota_project_id == "test-quota-project"
+
+ @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+ def test_quota_project_id_from_connection(self,
mock_get_creds_and_proj_id):
+ """Test that quota project ID from connection is applied to
credentials."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_get_creds_and_proj_id.return_value = (mock_creds, "test-project")
+
+ # Mock connection with quota_project_id in extras
+ uri = "google-cloud-platform://?quota_project_id=test-quota-project"
+ with mock.patch.dict("os.environ",
{"AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT": uri}):
+ hook = GoogleBaseHook(gcp_conn_id="google_cloud_default")
+ creds, _ = hook.get_credentials_and_project_id()
+
mock_creds.with_quota_project.assert_called_once_with("test-quota-project")
Review Comment:
nit
```suggestion
assert mock_creds.with_quota_project.mock_calls =
[call("test-quota-project")]
```
--
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]
