formatting

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -164,7 +164,9 @@ def __init__(
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
-        self._batch_size = batch_size if batch_size is not None else self.DEFAULT_BATCH_SIZE
+        self._batch_size = (
+            batch_size if batch_size is not None else self.DEFAULT_BATCH_SIZE
+        )
         self._flush_interval_seconds = self.DEFAULT_FLUSH_INTERVAL_SECONDS
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
@@ -175,26 +177,30 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
         self._flush_timer = None
-        
+
         # Start the periodic flush timer
         self._start_flush_timer()
 
     def _start_flush_timer(self):
         """Start the periodic flush timer"""
-        
+
         self._flush_timer = threading.Timer(
-            self._flush_interval_seconds, 
-            self._periodic_flush
+            self._flush_interval_seconds, self._periodic_flush
         )
         self._flush_timer.daemon = True  # Don't prevent program exit
         self._flush_timer.start()
-        logger.debug("Started flush timer for connection %s (interval: %d seconds)", 
-                    self._session_id_hex, self._flush_interval_seconds)
+        logger.debug(
+            "Started flush timer for connection %s (interval: %d seconds)",
+            self._session_id_hex,
+            self._flush_interval_seconds,
+        )
 
     def _periodic_flush(self):
         """Periodic flush callback - flushes events and reschedules the timer"""
-       
-        logger.debug("Performing periodic flush for connection %s", self._session_id_hex)
+
+        logger.debug(
+            "Performing periodic flush for connection %s", self._session_id_hex
+        )
         self._flush()
         # Reschedule the next flush
         self._start_flush_timer()
@@ -208,7 +214,7 @@ def _stop_flush_timer(self):
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
-            
+
         logger.debug("Exporting event for connection %s", self._session_id_hex)
         with self._lock:
             self._events_batch.append(event)
@@ -220,7 +226,7 @@ def _export_event(self, event):
 
     def _flush(self):
         """Flush the current batch of events to the server"""
-            
+
         with self._lock:
             events_to_flush = self._events_batch.copy()
             self._events_batch = []

tests/unit/test_telemetry.py

@@ -186,13 +186,15 @@ def test_batch_size_flush(self, telemetry_client_setup):
         client = telemetry_client_setup["client"]
         client._flush = MagicMock()
 
-        for i in range(TelemetryClient._batch_size-1):
+        batch_size = client._batch_size
+        
+        for i in range(batch_size - 1):
             client._export_event(f"event-{i}")
 
         client._flush.assert_not_called()
-        assert len(client._events_batch) == TelemetryClient._batch_size-1
+        assert len(client._events_batch) == batch_size - 1
 
-        client._export_event(f"event-{TelemetryClient._batch_size - 1}")
+        client._export_event(f"event-{batch_size - 1}")
 
         client._flush.assert_called_once()