diff --git a/src/snowflake/snowpark/mock/_telemetry.py b/src/snowflake/snowpark/mock/_telemetry.py index 39ebd9232..857291b47 100644 --- a/src/snowflake/snowpark/mock/_telemetry.py +++ b/src/snowflake/snowpark/mock/_telemetry.py @@ -91,6 +91,7 @@ def __init__(self) -> None: os.getenv("SNOWPARK_LOCAL_TESTING_INTERNAL_TELEMETRY", False) ) self._deployment_url = self.PROD + self._enable = True def _upload_payload(self, payload) -> None: if not REQUESTS_AVAILABLE: @@ -142,6 +143,19 @@ def add(self, event) -> None: return self._upload_payload(payload) + @property + def enabled(self) -> bool: + """Whether the Telemetry service is enabled or not.""" + return self._enabled + + def enable(self) -> None: + """Enable Telemetry Service.""" + self._enabled = True + + def disable(self) -> None: + """Disable Telemetry Service.""" + self._enabled = False + def export_queue_to_string(self): logs = list() while not self.queue.empty(): diff --git a/tests/mock/test_oob_telemetry.py b/tests/mock/test_oob_telemetry.py index f3e3d80f7..7016dc563 100644 --- a/tests/mock/test_oob_telemetry.py +++ b/tests/mock/test_oob_telemetry.py @@ -220,7 +220,16 @@ def test_unit_connection(caplog, local_testing_telemetry_setup): def test_unit_connection_disable_telemetry(caplog, local_testing_telemetry_setup): # clean up queue first - LocalTestOOBTelemetryService.get_instance().export_queue_to_string() + telemetry_service = LocalTestOOBTelemetryService.get_instance() + assert telemetry_service.enabled is True + telemetry_service.export_queue_to_string() + + telemetry_service.disable() + assert not telemetry_service.enabled + + telemetry_service.enable() + assert telemetry_service.enabled is True + disabled_telemetry_conn = MockServerConnection( options={"disable_local_testing_telemetry": True} )