Fix socket timeout test

Signed-off-by: Matthew Kim <[email protected]>

Author: mattdeekay

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 2.6.x (Unreleased)
 
+- Fix: test_socket_timeout_user_defined e2e test
+
 ## 2.6.1 (2023-06-08)
 
 - Fix: http.client would raise a BadStatusLine exception in some cases

tests/e2e/driver_tests.py

@@ -16,6 +16,7 @@
 import pytz
 import thrift
 import pytest
+from urllib3.connectionpool import ReadTimeoutError
 
 import databricks.sql as sql
 from databricks.sql import STRING, BINARY, NUMBER, DATETIME, DATE, DatabaseError, Error, OperationalError, RequestError
@@ -477,7 +478,7 @@ def test_socket_timeout_user_defined(self):
                 query = "select * from range(10000000)"
                 cursor.execute(query)
 
-        self.assertIsInstance(cm.exception.args[1], TimeoutError)
+        self.assertIsInstance(cm.exception.args[1], ReadTimeoutError)
 
     def test_ssp_passthrough(self):
         for enable_ansi in (True, False):