From 5eb45041dea5a08b30d7cd3afb168c37bea8d123 Mon Sep 17 00:00:00 2001 From: Rakshil Modi Date: Mon, 20 Jul 2026 11:54:44 -0700 Subject: [PATCH] retrive api changes --- awsiot/mqtt5_client_builder.py | 6 ++++- awsiot/mqtt_connection_builder.py | 5 +++- test/test_get_metrics.py | 44 ++++++++++++++++++++++++++----- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/awsiot/mqtt5_client_builder.py b/awsiot/mqtt5_client_builder.py index 99468c5c..95833635 100644 --- a/awsiot/mqtt5_client_builder.py +++ b/awsiot/mqtt5_client_builder.py @@ -170,6 +170,7 @@ **cipher_pref** (:class:`awscrt.io.TlsCipherPref`): Cipher preference to use for TLS connection. Default is `TlsCipherPref.DEFAULT`. + **enable_metrics_collection** (`bool`): Whether to send the SDK version number in the CONNECT packet. """ @@ -334,7 +335,10 @@ def _builder( tls_ctx = awscrt.io.ClientTlsContext(tls_ctx_options) client_options.tls_ctx = tls_ctx - client_options.metrics = _build_sdk_metrics() + if _get(kwargs, 'enable_metrics_collection', True): + client_options.metrics = _build_sdk_metrics() + else: + client_options.disable_metrics = True client = awscrt.mqtt5.Client(client_options=client_options) return client diff --git a/awsiot/mqtt_connection_builder.py b/awsiot/mqtt_connection_builder.py index d1f90db3..37aafdab 100644 --- a/awsiot/mqtt_connection_builder.py +++ b/awsiot/mqtt_connection_builder.py @@ -113,6 +113,8 @@ **cipher_pref** (:class:`awscrt.io.TlsCipherPref`): Cipher preference to use for TLS connection. Default is `TlsCipherPref.DEFAULT`. + **enable_metrics_collection** (`bool`): Whether to send the SDK version number in the CONNECT packet. + **http_proxy_options** (:class: 'awscrt.http.HttpProxyOptions'): HTTP proxy options to use """ @@ -199,7 +201,7 @@ def _builder( if username == "": username = None - metrics = _build_sdk_metrics() + metrics = _build_sdk_metrics() if _get(kwargs, 'enable_metrics_collection', True) else None client_bootstrap = _get(kwargs, 'client_bootstrap') if client_bootstrap is None: @@ -232,6 +234,7 @@ def _builder( on_connection_success=_get(kwargs, 'on_connection_success'), on_connection_failure=_get(kwargs, 'on_connection_failure'), on_connection_closed=_get(kwargs, 'on_connection_closed'), + disable_metrics=metrics is None, metrics=metrics ) diff --git a/test/test_get_metrics.py b/test/test_get_metrics.py index 41650b59..c098467f 100644 --- a/test/test_get_metrics.py +++ b/test/test_get_metrics.py @@ -10,6 +10,11 @@ import boto3 import botocore.exceptions +import awscrt.io +import awscrt.mqtt +import awscrt.mqtt5 + +from awsiot import mqtt_connection_builder, mqtt5_client_builder from awsiot._iot_metrics import ( _IOT_SDK_METRICS_VERSION, _get_sdk_version, @@ -108,9 +113,6 @@ class TestMqtt3BuilderMetrics(unittest.TestCase): def test_metrics_attached(self): """Builder should always pass SDK metrics to Connection.""" config = Config.get() - import awscrt.io - import awscrt.mqtt - from awsiot import mqtt_connection_builder with patch("awsiot._iot_metrics._get_sdk_version", return_value="2.0.0"), \ patch.object(awscrt.mqtt, "Connection") as mock_conn, \ @@ -127,16 +129,30 @@ def test_metrics_attached(self): self.assertEqual(entries["IoTSDKVersion"], "2.0.0") self.assertEqual(entries["IoTSDKMetricsVersion"], str(_IOT_SDK_METRICS_VERSION)) + def test_metrics_disabled_when_flag_false(self): + """Builder should suppress SDK metrics when enable_metrics_collection=False.""" + config = Config.get() + + with patch.object(awscrt.mqtt, "Connection") as mock_conn, \ + patch.object(awscrt.mqtt, "Client"): + mqtt_connection_builder._builder( + awscrt.io.TlsContextOptions(), + endpoint=config.endpoint, + client_id=create_client_id(), + enable_metrics_collection=False, + ) + + kwargs = mock_conn.call_args.kwargs + self.assertIsNone(kwargs["metrics"]) + self.assertTrue(kwargs["disable_metrics"]) + class TestMqtt5BuilderMetrics(unittest.TestCase): - """Test that mqtt5_client_builder attaches SDK metrics.""" + """Test that mqtt5_client_builder attaches SDK metrics.""" def test_metrics_attached(self): """Builder should set SDK metrics on client_options.""" config = Config.get() - import awscrt.io - import awscrt.mqtt5 - from awsiot import mqtt5_client_builder with patch("awsiot._iot_metrics._get_sdk_version", return_value="2.0.0"), \ patch.object(awscrt.mqtt5, "Client") as mock_client: @@ -151,6 +167,20 @@ def test_metrics_attached(self): self.assertEqual(entries["IoTSDKVersion"], "2.0.0") self.assertEqual(entries["IoTSDKMetricsVersion"], str(_IOT_SDK_METRICS_VERSION)) + def test_metrics_disabled_when_flag_false(self): + """Builder should suppress SDK metrics when enable_metrics_collection=False.""" + config = Config.get() + + with patch.object(awscrt.mqtt5, "Client") as mock_client: + mqtt5_client_builder._builder( + awscrt.io.TlsContextOptions(), + endpoint=config.endpoint, + enable_metrics_collection=False, + ) + + client_options = mock_client.call_args.kwargs["client_options"] + self.assertTrue(client_options.disable_metrics) + if __name__ == "__main__": unittest.main()