@@ -294,38 +294,4 @@ def test_factory_shutdown_flow(self):
294294 # Close second client - factory should shut down
295295 TelemetryClientFactory .close (session2 )
296296 assert TelemetryClientFactory ._initialized is False
297- assert TelemetryClientFactory ._executor is None
298-
299-
300- @patch ("databricks.sql.client.Session" )
301- def test_connection_with_custom_batch_size (MockSession ):
302- """
303- Verifies that creating a Connection with `telemetry_batch_size` results
304- in a TelemetryClient instance with the correct batch size attribute.
305- """
306-
307- custom_batch_size = 50
308- session_id = "session-with-custom-batch-size"
309-
310- # Mock the Session object created inside Connection to avoid network calls
311- mock_session_instance = MockSession .return_value
312- mock_session_instance .guid_hex = session_id
313- mock_session_instance .auth_provider = MagicMock ()
314- mock_session_instance .host = "mock-host"
315- mock_session_instance .port = 443
316- mock_session_instance .useragent_header = "mock-agent"
317- mock_session_instance .is_open = True
318-
319- from databricks .sql .client import Connection
320- conn = Connection (
321- server_hostname = "test-host" ,
322- http_path = "/test-path" ,
323- access_token = "test-token" ,
324- enable_telemetry = True ,
325- telemetry_batch_size = custom_batch_size
326- )
327-
328- client = TelemetryClientFactory .get_telemetry_client (session_id )
329- assert isinstance (client , TelemetryClient )
330- assert client ._batch_size == custom_batch_size
331- conn .close ()
297+ assert TelemetryClientFactory ._executor is None
0 commit comments