checking stuff

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

Author: saishreeeee

.github/workflows/code-quality-checks.yml

@@ -75,7 +75,7 @@ jobs:
           uses: actions/checkout@v2
       -   name: Set up python ${{ matrix.python-version }}
           id: setup-python
-          uses: actions/setup-python@v2
+          uses: actions/setup-python@v5
           with:
             python-version: ${{ matrix.python-version }}
       #----------------------------------------------

tests/unit/test_client.py

@@ -52,7 +52,6 @@ def new(cls):
         )
 
         ThriftBackendMock.execute_command.return_value = MockTExecuteStatementResp
-
         return ThriftBackendMock
 
     @classmethod
@@ -352,30 +351,30 @@ def test_context_manager_closes_cursor(self):
         finally:
             cursor.close.assert_called()
 
-    # @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
-    # def test_context_manager_closes_connection(self, mock_client_class):
-    #     print("hellow1")
-    #     instance = mock_client_class.return_value
-
-    #     mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
-    #     mock_open_session_resp.sessionHandle.sessionId = b"\x22"
-    #     instance.open_session.return_value = mock_open_session_resp
-
-    #     with databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) as connection:
-    #         pass
-
-    #     # Check the close session request has an id of x22
-    #     close_session_id = instance.close_session.call_args[0][0].sessionId
-    #     self.assertEqual(close_session_id, b"\x22")
-
-    #     connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
-    #     connection.close = Mock()
-    #     try:
-    #         with self.assertRaises(KeyboardInterrupt):
-    #             with connection:
-    #                 raise KeyboardInterrupt("Simulated interrupt")
-    #     finally:
-    #         connection.close.assert_called()
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_context_manager_closes_connection(self, mock_client_class):
+        print("hellow1")
+        instance = mock_client_class.return_value
+
+        mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
+        mock_open_session_resp.sessionHandle.sessionId = b"\x22"
+        instance.open_session.return_value = mock_open_session_resp
+
+        with databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) as connection:
+            pass
+
+        # Check the close session request has an id of x22
+        close_session_id = instance.close_session.call_args[0][0].sessionId
+        self.assertEqual(close_session_id, b"\x22")
+
+        connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
+        connection.close = Mock()
+        try:
+            with self.assertRaises(KeyboardInterrupt):
+                with connection:
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            connection.close.assert_called()
 
     def dict_product(self, dicts):
         """