removed walrus, removed random test before stalling test

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

Author: saishreeeee

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

@@ -112,7 +112,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       -   name: Run tests
-          run: poetry run python -m pytest tests/unit
+          run: poetry run python -m pytest tests/unit -v -s
   check-linting:
     runs-on: ubuntu-latest
     strategy:

src/databricks/sql/telemetry/telemetry_client.py

@@ -415,14 +415,16 @@ def close(session_id_hex):
         """Close and remove the telemetry client for a specific connection"""
 
         with TelemetryClientFactory._lock:
-            if (
-                telemetry_client := TelemetryClientFactory._clients.pop(
-                    session_id_hex, None
-                )
-            ) is not None:
+            # if (
+            #     telemetry_client := TelemetryClientFactory._clients.pop(
+            #         session_id_hex, None
+            #     )
+            # ) is not None:
+            if session_id_hex in TelemetryClientFactory._clients:
                 logger.debug(
                     "Removing telemetry client for connection %s", session_id_hex
                 )
+                telemetry_client = TelemetryClientFactory._clients.pop(session_id_hex)
                 telemetry_client.close()
 
             # Shutdown executor if no more clients

tests/unit/test_client.py

@@ -353,6 +353,7 @@ def test_context_manager_closes_cursor(self):
 
     @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
     def test_context_manager_closes_connection(self, mock_client_class):
+        print("stalling test")
         instance = mock_client_class.return_value
 
         mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
@@ -649,36 +650,36 @@ def test_disable_pandas_respected(self, mock_thrift_backend_class):
 
         mock_table.itercolumns.assert_called_once_with()
 
-    def test_column_name_api(self):
-        ResultRow = Row("first_col", "second_col", "third_col")
-        data = [
-            ResultRow("val1", 321, 52.32),
-            ResultRow("val2", 2321, 252.32),
-        ]
-
-        expected_values = [["val1", 321, 52.32], ["val2", 2321, 252.32]]
-
-        for (row, expected) in zip(data, expected_values):
-            self.assertEqual(row.first_col, expected[0])
-            self.assertEqual(row.second_col, expected[1])
-            self.assertEqual(row.third_col, expected[2])
-
-            self.assertEqual(row["first_col"], expected[0])
-            self.assertEqual(row["second_col"], expected[1])
-            self.assertEqual(row["third_col"], expected[2])
-
-            self.assertEqual(row[0], expected[0])
-            self.assertEqual(row[1], expected[1])
-            self.assertEqual(row[2], expected[2])
-
-            self.assertEqual(
-                row.asDict(),
-                {
-                    "first_col": expected[0],
-                    "second_col": expected[1],
-                    "third_col": expected[2],
-                },
-            )
+    # def test_column_name_api(self):
+    #     ResultRow = Row("first_col", "second_col", "third_col")
+    #     data = [
+    #         ResultRow("val1", 321, 52.32),
+    #         ResultRow("val2", 2321, 252.32),
+    #     ]
+
+    #     expected_values = [["val1", 321, 52.32], ["val2", 2321, 252.32]]
+
+    #     for (row, expected) in zip(data, expected_values):
+    #         self.assertEqual(row.first_col, expected[0])
+    #         self.assertEqual(row.second_col, expected[1])
+    #         self.assertEqual(row.third_col, expected[2])
+
+    #         self.assertEqual(row["first_col"], expected[0])
+    #         self.assertEqual(row["second_col"], expected[1])
+    #         self.assertEqual(row["third_col"], expected[2])
+
+    #         self.assertEqual(row[0], expected[0])
+    #         self.assertEqual(row[1], expected[1])
+    #         self.assertEqual(row[2], expected[2])
+
+    #         self.assertEqual(
+    #             row.asDict(),
+    #             {
+    #                 "first_col": expected[0],
+    #                 "second_col": expected[1],
+    #                 "third_col": expected[2],
+    #             },
+    #         )
 
     @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):