diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py index 3b1487c8c7bb..21278cd624bf 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py @@ -21,7 +21,7 @@ from google.cloud._storage_v2.services.storage.transports.base import ( DEFAULT_CLIENT_INFO, ) -from google.cloud.storage import __version__ +from google.cloud.storage import __version__, _opentelemetry_metrics from google.cloud.storage.grpc_client import ( _DEFAULT_HOST, _resolve_direct_path_interconnect, @@ -73,6 +73,19 @@ class AsyncGrpcClient: using xDS and standard TLS. Can also be overridden via the ``GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS_OVER_INTERCONNECT`` environment variable (``"true"`` or ``"false"``). Defaults to ``False``. + + :type enable_metrics: bool or None + :param enable_metrics: + (Optional, Experimental) Whether to enable OpenTelemetry metrics. If + ``None``, falls back to the ``GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS`` + environment variable, or ``False`` if unset. + + :type enable_debug_metrics: bool or None + :param enable_debug_metrics: + (Optional, Experimental) Whether to enable debug OpenTelemetry + metrics. If ``None``, falls back to the + ``GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS`` environment variable, + or ``False`` if unset. """ def __init__( @@ -83,7 +96,18 @@ def __init__( *, attempt_direct_path=True, attempt_direct_path_xds_over_interconnect=False, + enable_metrics=None, + enable_debug_metrics=None, ): + self._enable_metrics = enable_metrics + self._enable_debug_metrics = enable_debug_metrics + self._metrics_enabled = _opentelemetry_metrics.is_metrics_enabled( + enable_metrics + ) + self._debug_metrics_enabled = _opentelemetry_metrics.is_debug_metrics_enabled( + enable_debug_metrics + ) + if isinstance(credentials, auth_credentials.AnonymousCredentials): if client_options is None or client_options.api_endpoint is None: raise ValueError( @@ -115,6 +139,16 @@ def __init__( attempt_direct_path_xds_over_interconnect=use_dp_interconnect, ) + @property + def metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if OpenTelemetry metrics recording is active for this client.""" + return self._metrics_enabled + + @property + def debug_metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if debug metrics recording is active for this client.""" + return self._debug_metrics_enabled + def _create_anonymous_client(self, client_options, credentials): channel = grpc.aio.insecure_channel(client_options.api_endpoint) transport = storage_v2.services.storage.transports.StorageGrpcAsyncIOTransport( diff --git a/packages/google-cloud-storage/google/cloud/storage/client.py b/packages/google-cloud-storage/google/cloud/storage/client.py index 5559dbc7d763..f9d354a4d269 100644 --- a/packages/google-cloud-storage/google/cloud/storage/client.py +++ b/packages/google-cloud-storage/google/cloud/storage/client.py @@ -32,6 +32,7 @@ from google.cloud.client import ClientWithProject from google.cloud.exceptions import NotFound +from google.cloud.storage import _opentelemetry_metrics from google.cloud.storage._bucket_metadata_cache import BucketMetadataCache from google.cloud.storage._helpers import ( _DEFAULT_SCHEME, @@ -126,6 +127,19 @@ class Client(ClientWithProject): (Optional) An API key. Mutually exclusive with any other credentials. This parameter is an alias for setting `client_options.api_key` and will supercede any api key set in the `client_options` parameter. + + :type enable_metrics: bool or None + :param enable_metrics: + (Optional, Experimental) Whether to enable OpenTelemetry metrics. If + ``None``, falls back to the ``GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS`` + environment variable, or ``False`` if unset. + + :type enable_debug_metrics: bool or None + :param enable_debug_metrics: + (Optional, Experimental) Whether to enable debug OpenTelemetry + metrics. If ``None``, falls back to the + ``GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS`` environment variable, + or ``False`` if unset. """ SCOPE = ( @@ -146,6 +160,8 @@ def __init__( extra_headers={}, *, api_key=None, + enable_metrics=None, + enable_debug_metrics=None, ): self._base_connection = None @@ -163,6 +179,14 @@ def __init__( self._initial_client_info = client_info self._initial_client_options = client_options self._extra_headers = extra_headers + self._enable_metrics = enable_metrics + self._enable_debug_metrics = enable_debug_metrics + self._metrics_enabled = _opentelemetry_metrics.is_metrics_enabled( + enable_metrics + ) + self._debug_metrics_enabled = _opentelemetry_metrics.is_debug_metrics_enabled( + enable_debug_metrics + ) connection_kw_args = {"client_info": client_info} @@ -294,6 +318,16 @@ def __init__( self._batch_stack = _LocalStack() self._bucket_metadata_cache = BucketMetadataCache(self) + @property + def metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if OpenTelemetry metrics recording is active for this client.""" + return self._metrics_enabled + + @property + def debug_metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if debug metrics recording is active for this client.""" + return self._debug_metrics_enabled + def close(self): """Close the client and clear any cached metadata or active connections.""" if hasattr(self, "_bucket_metadata_cache") and self._bucket_metadata_cache: diff --git a/packages/google-cloud-storage/google/cloud/storage/grpc_client.py b/packages/google-cloud-storage/google/cloud/storage/grpc_client.py index b3d8b261e77f..79bcb4f0d8b1 100644 --- a/packages/google-cloud-storage/google/cloud/storage/grpc_client.py +++ b/packages/google-cloud-storage/google/cloud/storage/grpc_client.py @@ -19,6 +19,7 @@ from google.cloud.client import ClientWithProject from google.cloud import _storage_v2 as storage_v2 +from google.cloud.storage import _opentelemetry_metrics _marker = object() _DEFAULT_HOST = "storage.googleapis.com" @@ -103,6 +104,19 @@ class GrpcClient(ClientWithProject): using xDS and standard TLS. Can also be overridden via the ``GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS_OVER_INTERCONNECT`` environment variable (``"true"`` or ``"false"``). Defaults to ``False``. + + :type enable_metrics: bool or None + :param enable_metrics: + (Optional, Experimental) Whether to enable OpenTelemetry metrics. If + ``None``, falls back to the ``GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS`` + environment variable, or ``False`` if unset. + + :type enable_debug_metrics: bool or None + :param enable_debug_metrics: + (Optional, Experimental) Whether to enable debug OpenTelemetry + metrics. If ``None``, falls back to the + ``GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS`` environment variable, + or ``False`` if unset. """ def __init__( @@ -115,9 +129,20 @@ def __init__( api_key=None, attempt_direct_path=True, attempt_direct_path_xds_over_interconnect=False, + enable_metrics=None, + enable_debug_metrics=None, ): super(GrpcClient, self).__init__(project=project, credentials=credentials) + self._enable_metrics = enable_metrics + self._enable_debug_metrics = enable_debug_metrics + self._metrics_enabled = _opentelemetry_metrics.is_metrics_enabled( + enable_metrics + ) + self._debug_metrics_enabled = _opentelemetry_metrics.is_debug_metrics_enabled( + enable_debug_metrics + ) + if isinstance(client_options, dict): if api_key: client_options["api_key"] = api_key @@ -138,6 +163,16 @@ def __init__( attempt_direct_path_xds_over_interconnect=use_dp_interconnect, ) + @property + def metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if OpenTelemetry metrics recording is active for this client.""" + return self._metrics_enabled + + @property + def debug_metrics_enabled(self) -> bool: + """(Experimental) Returns ``True`` if debug metrics recording is active for this client.""" + return self._debug_metrics_enabled + def _create_gapic_client( self, credentials=None, diff --git a/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py b/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py index 6486696021fc..b376c5138d72 100644 --- a/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py +++ b/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py @@ -1383,6 +1383,8 @@ def _reduce_client(cl): client_info = cl._initial_client_info client_options = cl._initial_client_options extra_headers = getattr(cl, "_extra_headers", {}) + enable_metrics = getattr(cl, "_enable_metrics", None) + enable_debug_metrics = getattr(cl, "_enable_debug_metrics", None) return _LazyClient, ( client_object_id, @@ -1392,6 +1394,8 @@ def _reduce_client(cl): client_info, client_options, extra_headers, + enable_metrics, + enable_debug_metrics, ) @@ -1453,6 +1457,13 @@ def _digest_ordered_checksum_and_size_pairs(checksum_and_size_pairs): return crc_digest +# Number of positional Client constructor arguments passed after `id` by +# `_reduce_client`: (project, credentials, _http, client_info, client_options). +# Trailing items in `args` correspond to `extra_headers` and keyword-only +# `enable_metrics` and `enable_debug_metrics` parameters. +_REDUCE_CLIENT_POS_ARGS = 5 + + class _LazyClient: """An object that will transform into either a cached or a new Client""" @@ -1461,6 +1472,17 @@ def __new__(cls, id, *args, **kwargs): if cached_client: return cached_client else: + if len(args) > _REDUCE_CLIENT_POS_ARGS: + kwargs.setdefault("extra_headers", args[_REDUCE_CLIENT_POS_ARGS]) + if len(args) > _REDUCE_CLIENT_POS_ARGS + 1: + kwargs.setdefault( + "enable_metrics", args[_REDUCE_CLIENT_POS_ARGS + 1] + ) + if len(args) > _REDUCE_CLIENT_POS_ARGS + 2: + kwargs.setdefault( + "enable_debug_metrics", args[_REDUCE_CLIENT_POS_ARGS + 2] + ) + args = args[:_REDUCE_CLIENT_POS_ARGS] cached_client = Client(*args, **kwargs) _cached_clients[id] = cached_client return cached_client diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py index 9dffa6b36c86..0b96a92eb055 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py @@ -149,6 +149,107 @@ def test_constructor_disables_directpath(self, mock_async_storage_client): mock_channel = mock_transport_cls.create_channel.return_value mock_transport_cls.assert_called_once_with(channel=mock_channel) + @mock.patch("google.cloud._storage_v2.StorageAsyncClient") + def test_metrics_properties(self, mock_async_storage_client): + from google.cloud.storage import _opentelemetry_metrics + + mock_transport_cls = mock.MagicMock() + mock_async_storage_client.get_transport_class.return_value = mock_transport_cls + mock_creds = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + client = async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_metrics=True, + enable_debug_metrics=True, + ) + assert client.metrics_enabled is True + assert client.debug_metrics_enabled is True + + client_base_false = async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_metrics=False, + enable_debug_metrics=True, + ) + assert client_base_false.metrics_enabled is False + assert client_base_false.debug_metrics_enabled is True + + client_both_false = async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_metrics=False, + enable_debug_metrics=False, + ) + assert client_both_false.metrics_enabled is False + assert client_both_false.debug_metrics_enabled is False + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=False, + ): + client = async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_metrics=True, + enable_debug_metrics=True, + ) + assert client.metrics_enabled is False + assert client.debug_metrics_enabled is False + + @mock.patch("google.cloud._storage_v2.StorageAsyncClient") + def test_ctor_w_invalid_enable_metrics_raises_type_error( + self, mock_async_storage_client + ): + mock_transport_cls = mock.MagicMock() + mock_async_storage_client.get_transport_class.return_value = mock_transport_cls + mock_creds = _make_credentials() + + with pytest.raises(TypeError): + async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_metrics="true", + ) + with pytest.raises(TypeError): + async_grpc_client.AsyncGrpcClient( + credentials=mock_creds, + enable_debug_metrics="true", + ) + + @mock.patch("google.cloud._storage_v2.StorageAsyncClient") + def test_ctor_metrics_defaults_and_env_fallback(self, mock_async_storage_client): + from google.cloud.storage import _opentelemetry_metrics + + mock_transport_cls = mock.MagicMock() + mock_async_storage_client.get_transport_class.return_value = mock_transport_cls + mock_creds = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + with mock.patch.dict("os.environ", {}, clear=True): + client = async_grpc_client.AsyncGrpcClient(credentials=mock_creds) + assert client._enable_metrics is None + assert client._enable_debug_metrics is None + assert client.metrics_enabled is False + assert client.debug_metrics_enabled is False + + with mock.patch.dict( + "os.environ", + { + "GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS": "true", + "GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS": "true", + }, + clear=True, + ): + client = async_grpc_client.AsyncGrpcClient(credentials=mock_creds) + assert client.metrics_enabled is True + assert client.debug_metrics_enabled is True + @mock.patch("google.cloud._storage_v2.StorageAsyncClient") def test_grpc_client_property(self, mock_grpc_gapic_client): # Arrange diff --git a/packages/google-cloud-storage/tests/unit/test_client.py b/packages/google-cloud-storage/tests/unit/test_client.py index 87d82581490d..672bde151053 100644 --- a/packages/google-cloud-storage/tests/unit/test_client.py +++ b/packages/google-cloud-storage/tests/unit/test_client.py @@ -259,6 +259,139 @@ def test_ctor_w_universe_domain_and_matched_credentials(self): self.assertEqual(client.api_endpoint, expected_api_endpoint) self.assertEqual(client.universe_domain, universe_domain) + def test_ctor_w_enable_metrics(self): + PROJECT = "PROJECT" + credentials = _make_credentials() + client = self._make_one( + project=PROJECT, credentials=credentials, enable_metrics=True + ) + self.assertTrue(client._enable_metrics) + + def test_ctor_w_enable_debug_metrics(self): + PROJECT = "PROJECT" + credentials = _make_credentials() + client = self._make_one( + project=PROJECT, credentials=credentials, enable_debug_metrics=True + ) + self.assertTrue(client._enable_debug_metrics) + + def test_ctor_w_invalid_enable_metrics_raises_type_error(self): + PROJECT = "PROJECT" + credentials = _make_credentials() + with self.assertRaises(TypeError): + self._make_one( + project=PROJECT, credentials=credentials, enable_metrics="true" + ) + with self.assertRaises(TypeError): + self._make_one( + project=PROJECT, + credentials=credentials, + enable_debug_metrics="true", + ) + + def test_client_metrics_enabled_property(self): + from google.cloud.storage import _opentelemetry_metrics + + PROJECT = "PROJECT" + credentials = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + client = self._make_one( + project=PROJECT, credentials=credentials, enable_metrics=True + ) + self.assertTrue(client.metrics_enabled) + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=False, + ): + client = self._make_one( + project=PROJECT, credentials=credentials, enable_metrics=True + ) + self.assertFalse(client.metrics_enabled) + + def test_client_debug_metrics_enabled_property(self): + from google.cloud.storage import _opentelemetry_metrics + + PROJECT = "PROJECT" + credentials = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + client = self._make_one( + project=PROJECT, + credentials=credentials, + enable_metrics=True, + enable_debug_metrics=True, + ) + self.assertTrue(client.debug_metrics_enabled) + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=False, + ): + client = self._make_one( + project=PROJECT, + credentials=credentials, + enable_metrics=True, + enable_debug_metrics=True, + ) + self.assertFalse(client.debug_metrics_enabled) + + # Verify that debug metrics are independent when base metrics are disabled + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + client_disabled_base = self._make_one( + project=PROJECT, + credentials=credentials, + enable_metrics=False, + enable_debug_metrics=True, + ) + self.assertFalse(client_disabled_base.metrics_enabled) + self.assertTrue(client_disabled_base.debug_metrics_enabled) + + def test_ctor_metrics_defaults_and_env_fallback(self): + from google.cloud.storage import _opentelemetry_metrics + + PROJECT = "PROJECT" + credentials = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + with mock.patch.dict("os.environ", {}, clear=True): + client = self._make_one(project=PROJECT, credentials=credentials) + self.assertIsNone(client._enable_metrics) + self.assertIsNone(client._enable_debug_metrics) + self.assertFalse(client.metrics_enabled) + self.assertFalse(client.debug_metrics_enabled) + + with mock.patch.dict( + "os.environ", + { + "GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS": "true", + "GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS": "true", + }, + clear=True, + ): + client = self._make_one(project=PROJECT, credentials=credentials) + self.assertTrue(client.metrics_enabled) + self.assertTrue(client.debug_metrics_enabled) + def test_ctor_w_universe_domain_and_mismatched_credentials(self): PROJECT = "PROJECT" universe_domain = "example.com" diff --git a/packages/google-cloud-storage/tests/unit/test_grpc_client.py b/packages/google-cloud-storage/tests/unit/test_grpc_client.py index 73a8e5b273e4..5e3789549696 100644 --- a/packages/google-cloud-storage/tests/unit/test_grpc_client.py +++ b/packages/google-cloud-storage/tests/unit/test_grpc_client.py @@ -360,3 +360,107 @@ def test_rewrite_host_for_interconnect_delimiters(self): grpc_client._rewrite_host_for_interconnect(original_host), expected_host, ) + + @mock.patch("google.cloud._storage_v2.StorageClient") + def test_metrics_properties(self, mock_storage_client): + from google.cloud.storage import _opentelemetry_metrics + + mock_creds = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + client = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_metrics=True, + enable_debug_metrics=True, + ) + self.assertTrue(client.metrics_enabled) + self.assertTrue(client.debug_metrics_enabled) + + client_base_false = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_metrics=False, + enable_debug_metrics=True, + ) + self.assertFalse(client_base_false.metrics_enabled) + self.assertTrue(client_base_false.debug_metrics_enabled) + + client_both_false = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_metrics=False, + enable_debug_metrics=False, + ) + self.assertFalse(client_both_false.metrics_enabled) + self.assertFalse(client_both_false.debug_metrics_enabled) + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=False, + ): + client = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_metrics=True, + enable_debug_metrics=True, + ) + self.assertFalse(client.metrics_enabled) + self.assertFalse(client.debug_metrics_enabled) + + @mock.patch("google.cloud._storage_v2.StorageClient") + def test_ctor_w_invalid_enable_metrics_raises_type_error(self, mock_storage_client): + mock_creds = _make_credentials() + with self.assertRaises(TypeError): + grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_metrics="true", + ) + with self.assertRaises(TypeError): + grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + enable_debug_metrics="true", + ) + + @mock.patch("google.cloud._storage_v2.StorageClient") + def test_ctor_metrics_defaults_and_env_fallback(self, mock_storage_client): + from google.cloud.storage import _opentelemetry_metrics + + mock_creds = _make_credentials() + + with mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ): + with mock.patch.dict("os.environ", {}, clear=True): + client = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + ) + self.assertIsNone(client._enable_metrics) + self.assertIsNone(client._enable_debug_metrics) + self.assertFalse(client.metrics_enabled) + self.assertFalse(client.debug_metrics_enabled) + + with mock.patch.dict( + "os.environ", + { + "GCP_STORAGE_PYTHON_ENABLE_OTEL_METRICS": "true", + "GCP_STORAGE_PYTHON_ENABLE_OTEL_DEBUG_METRICS": "true", + }, + clear=True, + ): + client = grpc_client.GrpcClient( + project="test-project", + credentials=mock_creds, + ) + self.assertTrue(client.metrics_enabled) + self.assertTrue(client.debug_metrics_enabled) diff --git a/packages/google-cloud-storage/tests/unit/test_transfer_manager.py b/packages/google-cloud-storage/tests/unit/test_transfer_manager.py index b9d6f11a742d..c370f2eede2a 100644 --- a/packages/google-cloud-storage/tests/unit/test_transfer_manager.py +++ b/packages/google-cloud-storage/tests/unit/test_transfer_manager.py @@ -1405,16 +1405,45 @@ def test__reduce_client(): "x-goog-custom-audit-foo": "bar", } client._extra_headers = custom_headers + client._enable_metrics = True + client._enable_debug_metrics = True with ( mock.patch( "google.cloud.storage.transfer_manager._cached_clients", new=fake_cache ), - mock.patch("google.cloud.storage.transfer_manager.Client"), + mock.patch("google.cloud.storage.transfer_manager.Client") as mock_client_cls, ): - replicated_client, kwargs = transfer_manager._reduce_client(client) + replicated_client, args = transfer_manager._reduce_client(client) assert replicated_client is not None - assert custom_headers in kwargs + assert custom_headers in args + assert args[-2:] == (True, True) + replicated_client(*args) + mock_client_cls.assert_called_once_with( + *args[1:6], + extra_headers=custom_headers, + enable_metrics=True, + enable_debug_metrics=True, + ) + + # Verify compatibility when only enable_metrics is present (len(args) == 7 after id) + fake_cache.clear() + mock_client_cls.reset_mock() + replicated_client(*args[:-1]) + mock_client_cls.assert_called_once_with( + *args[1:6], + extra_headers=custom_headers, + enable_metrics=True, + ) + + # Verify compatibility when only extra_headers is present (len(args) == 6 after id) + fake_cache.clear() + mock_client_cls.reset_mock() + replicated_client(*args[:-2]) + mock_client_cls.assert_called_once_with( + *args[1:6], + extra_headers=custom_headers, + ) def test__call_method_on_maybe_pickled_blob():