test(telemetry/e2e): make TestTelemetryE2E deterministic by capturing submissions, not callbacks (#812)

test(telemetry/e2e): make TestTelemetryE2E deterministic

The previous tests asserted "telemetry round-tripped to the server"
by intercepting TelemetryClient._telemetry_request_callback and
counting completed futures. That recording lags the actual work —
the callback fires asynchronously after the HTTP request completes,
and on the *last* connection close TelemetryClientFactory.close()
shuts the shared executor down with wait=False (intentional, for
connection-close latency in production). Two consequences:

  1. A `wait(captured_futures, timeout=10)` call right after `with
     conn:` can return before any callbacks have fired — so the
     wait is "waiting on" an empty list, returns immediately, and
     the assertion `assert len(done) == expected_count` fails non-
     deterministically with `assert 0 == 2` or `assert 1 == 2`.

  2. The shared-executor shutdown(wait=False) can drop in-flight
     submissions that haven't started running yet, so even if we
     drained correctly we'd be testing whether the server happened
     to receive the request in time, not whether the connector
     correctly dispatched it.

Switch interception from `_telemetry_request_callback` to
`_send_telemetry`. That captures the connector's *intent to submit*
synchronously, which is what we actually want to test — the
connector either decided to send a batch or it didn't, regardless
of what happens to the future afterward. No sleep needed, no
timeout-based wait needed, no race against the executor shutdown.

5 consecutive local runs pass deterministically in ~20s each (down
from ~17 min when the flake hit).

Co-authored-by: Isaac

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

tests/e2e/test_telemetry_e2e.py

@@ -1,13 +1,11 @@
 """
 E2E test for telemetry - verifies telemetry behavior with different scenarios
 """
-import time
 import threading
 import logging
 from contextlib import contextmanager
 from unittest.mock import patch
 import pytest
-from concurrent.futures import wait
 
 import databricks.sql as sql
 from databricks.sql.telemetry.telemetry_client import (
@@ -91,25 +89,43 @@ def telemetry_setup_teardown(self):
 
     @pytest.fixture
     def telemetry_interceptors(self):
-        """Setup reusable telemetry interceptors as a fixture"""
+        """Setup reusable telemetry interceptors as a fixture.
+
+        Captures two signals:
+          * captured_events — every call to TelemetryClient._export_event
+            (the connector intended to send this event).
+          * captured_submissions — every call to TelemetryClient._send_telemetry
+            (the connector submitted this batch to the executor).
+
+        We deliberately do NOT capture inside _telemetry_request_callback. That
+        callback only fires after the HTTP round-trip completes, and the
+        connector's TelemetryClientFactory.close() shuts down the shared
+        executor with wait=False — so on the last connection close, in-flight
+        futures may never get to run their callbacks before the executor is
+        gone. That's intentional for connection-close latency in production,
+        but it means a callback-time assertion is racing the connector's
+        shutdown path. Capturing at submission time tests what the connector
+        actually controls.
+        """
         capture_lock = threading.Lock()
         captured_events = []
-        captured_futures = []
+        captured_submissions = []
 
         original_export = TelemetryClient._export_event
-        original_callback = TelemetryClient._telemetry_request_callback
+        original_send = TelemetryClient._send_telemetry
 
         def export_wrapper(self_client, event):
             with capture_lock:
                 captured_events.append(event)
             return original_export(self_client, event)
 
-        def callback_wrapper(self_client, future, sent_count):
+        def send_wrapper(self_client, events):
             with capture_lock:
-                captured_futures.append(future)
-            original_callback(self_client, future, sent_count)
+                # Record the batch (list of events) the connector submitted.
+                captured_submissions.append(list(events))
+            return original_send(self_client, events)
 
-        return captured_events, captured_futures, export_wrapper, callback_wrapper
+        return captured_events, captured_submissions, export_wrapper, send_wrapper
 
     # ==================== ASSERTION HELPERS ====================
 
@@ -165,24 +181,33 @@ def assert_error_info(self, event, expected_error_name=None):
         if expected_error_name:
             assert error_info.error_name == expected_error_name
 
-    def verify_events(self, captured_events, captured_futures, expected_count):
-        """Common verification for event count and HTTP responses"""
+    def verify_events(self, captured_events, captured_submissions, expected_count):
+        """Common verification for event count and submission count.
+
+        Asserts on what the connector did — exported events and submitted
+        batches — not on what the server returned. Server-side HTTP success
+        is asserted via end-to-end behavior elsewhere (and would race the
+        connector's wait=False executor shutdown on connection close, see
+        the docstring on telemetry_interceptors).
+
+        Because these tests use telemetry_batch_size=1, each exported event
+        triggers its own batch submission, so the submission count should
+        equal the event count.
+        """
         if expected_count == 0:
-            assert len(captured_events) == 0, f"Expected 0 events, got {len(captured_events)}"
-            assert len(captured_futures) == 0, f"Expected 0 responses, got {len(captured_futures)}"
+            assert len(captured_events) == 0, \
+                f"Expected 0 events, got {len(captured_events)}"
+            assert len(captured_submissions) == 0, \
+                f"Expected 0 submissions, got {len(captured_submissions)}"
         else:
             assert len(captured_events) == expected_count, \
                 f"Expected {expected_count} events, got {len(captured_events)}"
+            # batch_size=1, so one submission per event.
+            submitted_event_count = sum(len(batch) for batch in captured_submissions)
+            assert submitted_event_count == expected_count, \
+                f"Expected {expected_count} submitted events across batches, " \
+                f"got {submitted_event_count} (batches: {[len(b) for b in captured_submissions]})"
 
-            time.sleep(2)
-            done, _ = wait(captured_futures, timeout=10)
-            assert len(done) == expected_count, \
-                f"Expected {expected_count} responses, got {len(done)}"
-            
-            for future in done:
-                response = future.result()
-                assert 200 <= response.status < 300
-            
             # Assert common fields for all events
             for event in captured_events:
                 self.assert_system_config(event)
@@ -199,11 +224,11 @@ def verify_events(self, captured_events, captured_futures, expected_count):
     def test_telemetry_flags(self, telemetry_interceptors, enable_telemetry, 
                             force_enable, expected_count, test_id):
         """Test telemetry behavior with different flag combinations"""
-        captured_events, captured_futures, export_wrapper, callback_wrapper = \
+        captured_events, captured_submissions, export_wrapper, send_wrapper = \
             telemetry_interceptors
 
         with patch.object(TelemetryClient, "_export_event", export_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+             patch.object(TelemetryClient, "_send_telemetry", send_wrapper):
 
             extra_params = {"telemetry_batch_size": 1}
             if enable_telemetry is not None:
@@ -216,9 +241,7 @@ def test_telemetry_flags(self, telemetry_interceptors, enable_telemetry,
                     cursor.execute("SELECT 1")
                     cursor.fetchone()
 
-            # Give time for async telemetry submission after connection closes
-            time.sleep(0.5)
-            self.verify_events(captured_events, captured_futures, expected_count)
+            self.verify_events(captured_events, captured_submissions, expected_count)
 
             # Assert statement execution on latency event (if events exist)
             if expected_count > 0:
@@ -230,11 +253,11 @@ def test_telemetry_flags(self, telemetry_interceptors, enable_telemetry,
     ])
     def test_sql_errors(self, telemetry_interceptors, query, expected_error):
         """Test telemetry captures error information for different SQL errors"""
-        captured_events, captured_futures, export_wrapper, callback_wrapper = \
+        captured_events, captured_submissions, export_wrapper, send_wrapper = \
             telemetry_interceptors
 
         with patch.object(TelemetryClient, "_export_event", export_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+             patch.object(TelemetryClient, "_send_telemetry", send_wrapper):
 
             with self.connection(extra_params={
                 "force_enable_telemetry": True,
@@ -245,9 +268,6 @@ def test_sql_errors(self, telemetry_interceptors, query, expected_error):
                         cursor.execute(query)
                         cursor.fetchone()
 
-            time.sleep(2)
-            wait(captured_futures, timeout=10)
-
             assert len(captured_events) >= 1
 
             # Find event with error_info
@@ -261,11 +281,11 @@ def test_sql_errors(self, telemetry_interceptors, query, expected_error):
 
     def test_metadata_operation(self, telemetry_interceptors):
         """Test telemetry for metadata operations (getCatalogs)"""
-        captured_events, captured_futures, export_wrapper, callback_wrapper = \
+        captured_events, captured_submissions, export_wrapper, send_wrapper = \
             telemetry_interceptors
 
         with patch.object(TelemetryClient, "_export_event", export_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+             patch.object(TelemetryClient, "_send_telemetry", send_wrapper):
 
             with self.connection(extra_params={
                 "force_enable_telemetry": True,
@@ -275,21 +295,18 @@ def test_metadata_operation(self, telemetry_interceptors):
                     catalogs = cursor.catalogs()
                     catalogs.fetchall()
 
-            time.sleep(2)
-            wait(captured_futures, timeout=10)
-
             assert len(captured_events) >= 1
             for event in captured_events:
                 self.assert_system_config(event)
                 self.assert_connection_params(event, self.arguments["http_path"])
 
     def test_direct_results(self, telemetry_interceptors):
         """Test telemetry with direct results (use_cloud_fetch=False)"""
-        captured_events, captured_futures, export_wrapper, callback_wrapper = \
+        captured_events, captured_submissions, export_wrapper, send_wrapper = \
             telemetry_interceptors
 
         with patch.object(TelemetryClient, "_export_event", export_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+             patch.object(TelemetryClient, "_send_telemetry", send_wrapper):
 
             with self.connection(extra_params={
                 "force_enable_telemetry": True,
@@ -301,9 +318,6 @@ def test_direct_results(self, telemetry_interceptors):
                     result = cursor.fetchall()
                     assert len(result) == 1 and result[0][0] == 100
 
-            time.sleep(2)
-            wait(captured_futures, timeout=10)
-
             assert len(captured_events) >= 2
             for event in captured_events:
                 self.assert_system_config(event)
@@ -320,11 +334,11 @@ def test_direct_results(self, telemetry_interceptors):
     def test_cloudfetch_with_different_close_patterns(self, telemetry_interceptors, 
                                                        close_type):
         """Test telemetry with cloud fetch using different resource closing patterns"""
-        captured_events, captured_futures, export_wrapper, callback_wrapper = \
+        captured_events, captured_submissions, export_wrapper, send_wrapper = \
             telemetry_interceptors
 
         with patch.object(TelemetryClient, "_export_event", export_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+             patch.object(TelemetryClient, "_send_telemetry", send_wrapper):
 
             if close_type == "explicit_connection":
                 # Test explicit connection close
@@ -365,9 +379,6 @@ def test_cloudfetch_with_different_close_patterns(self, telemetry_interceptors,
                         result = cursor.fetchall()
                         assert len(result) == 1000
 
-            time.sleep(2)
-            wait(captured_futures, timeout=10)
-
             assert len(captured_events) >= 2
             for event in captured_events:
                 self.assert_system_config(event)