From af634b0020f09da036aeb93a67e1261fae9e7d88 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Fri, 7 Aug 2026 23:43:33 +0000 Subject: [PATCH] fix(storage): mock wrap_method to verify client_info in tests --- .../tests/unit/test_read_client_v1.py | 26 ++++++------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py index dd29d6c19fbc..c76700dc2ffc 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py @@ -19,7 +19,6 @@ import pytest from google.api_core.gapic_v1 import client_info from google.auth import credentials - from google.cloud.bigquery_storage_v1 import types PROJECT = "my-project" @@ -158,24 +157,15 @@ def test_read_rows(mock_transport, client_under_test): ["google.cloud.bigquery_storage_v1", "google.cloud.bigquery_storage_v1beta2"], ) def test_init_default_client_info(module_under_test): - from google.api_core.gapic_v1.client_info import METRICS_METADATA_KEY - mut = importlib.import_module(module_under_test) - creds = mock.Mock(spec=credentials.Credentials) - client = mut.BigQueryWriteClient(credentials=creds) + expected_client_info = f"gccl/{mut.__version__}" - installed_version = mut.__version__ - expected_client_info = f"gccl/{installed_version}" + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as mock_wrap: + mut.BigQueryWriteClient(credentials=creds) - for wrapped_method in client.transport._wrapped_methods.values(): - user_agent = next( - ( - header_value - for header, header_value in wrapped_method._metadata - if header == METRICS_METADATA_KEY - ), - None, - ) - assert user_agent is not None - assert expected_client_info in user_agent + assert mock_wrap.call_count > 0 + for call in mock_wrap.call_args_list: + client_info = call.kwargs.get("client_info") + assert client_info is not None + assert expected_client_info in client_info.to_user_agent()