From ab0b76adf5a3096914df0ed4036068e099b876a3 Mon Sep 17 00:00:00 2001
From: Adam Ling <adam.ling@snowflake.com>
Date: Wed, 3 Jul 2024 15:49:38 -0700
Subject: [PATCH 1/2] poc

---
 src/snowflake/snowpark/mock/_telemetry.py | 30 +++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/snowflake/snowpark/mock/_telemetry.py b/src/snowflake/snowpark/mock/_telemetry.py
index 39ebd9232..2ed7dd67c 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():
@@ -224,5 +238,21 @@ def log_not_supported_error(
                     raise_error(error_message), error_message
                 )
 
+    def close(self) -> None:
+        """Closes the telemetry service."""
+        self.flush()
+        self.disable()
+
+    def flush(self) -> None:
+        """Flushes all telemetry events in the queue and submit them to the back-end."""
+        if not self.enabled:
+            return
+
+        if not self.queue.empty():
+            payload = self.export_queue_to_string()
+            if payload is None:
+                return
+            self._upload_payload(payload)
+
 
 atexit.register(LocalTestOOBTelemetryService.get_instance().close)

From d16dee166775d023b32d2452d7fb13ae7c923ee3 Mon Sep 17 00:00:00 2001
From: Adam Ling <adam.ling@snowflake.com>
Date: Thu, 11 Jul 2024 10:44:53 -0700
Subject: [PATCH 2/2] overwrite connector methods

---
 src/snowflake/snowpark/mock/_telemetry.py | 16 ----------------
 tests/mock/test_oob_telemetry.py          | 11 ++++++++++-
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/src/snowflake/snowpark/mock/_telemetry.py b/src/snowflake/snowpark/mock/_telemetry.py
index 2ed7dd67c..857291b47 100644
--- a/src/snowflake/snowpark/mock/_telemetry.py
+++ b/src/snowflake/snowpark/mock/_telemetry.py
@@ -238,21 +238,5 @@ def log_not_supported_error(
                     raise_error(error_message), error_message
                 )
 
-    def close(self) -> None:
-        """Closes the telemetry service."""
-        self.flush()
-        self.disable()
-
-    def flush(self) -> None:
-        """Flushes all telemetry events in the queue and submit them to the back-end."""
-        if not self.enabled:
-            return
-
-        if not self.queue.empty():
-            payload = self.export_queue_to_string()
-            if payload is None:
-                return
-            self._upload_payload(payload)
-
 
 atexit.register(LocalTestOOBTelemetryService.get_instance().close)
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}
     )