align test_session with pytest instead of unittest

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

tests/unit/test_session.py

@@ -1,4 +1,4 @@
-import unittest
+import pytest
 from unittest.mock import patch, MagicMock, Mock, PropertyMock
 import gc
 
@@ -12,7 +12,7 @@
 import databricks.sql
 
 
-class SessionTestSuite(unittest.TestCase):
+class TestSession:
     """
     Unit tests for Session functionality
     """
@@ -37,8 +37,8 @@ def test_close_uses_the_correct_session_id(self, mock_client_class):
 
         # Check that close_session was called with the correct SessionId
         close_session_call_args = instance.close_session.call_args[0][0]
-        self.assertEqual(close_session_call_args.guid, b"\x22")
-        self.assertEqual(close_session_call_args.secret, b"\x33")
+        assert close_session_call_args.guid == b"\x22"
+        assert close_session_call_args.secret == b"\x33"
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_auth_args(self, mock_client_class):
@@ -63,8 +63,8 @@ def test_auth_args(self, mock_client_class):
         for args in connection_args:
             connection = databricks.sql.connect(**args)
             host, port, http_path, *_ = mock_client_class.call_args[0]
-            self.assertEqual(args["server_hostname"], host)
-            self.assertEqual(args["http_path"], http_path)
+            assert args["server_hostname"] == host
+            assert args["http_path"] == http_path
             connection.close()
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
@@ -73,7 +73,7 @@ def test_http_header_passthrough(self, mock_client_class):
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, http_headers=http_headers)
 
         call_args = mock_client_class.call_args[0][3]
-        self.assertIn(("foo", "bar"), call_args)
+        assert ("foo", "bar") in call_args
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_tls_arg_passthrough(self, mock_client_class):
@@ -86,10 +86,10 @@ def test_tls_arg_passthrough(self, mock_client_class):
         )
 
         kwargs = mock_client_class.call_args[1]
-        self.assertEqual(kwargs["_tls_verify_hostname"], "hostname")
-        self.assertEqual(kwargs["_tls_trusted_ca_file"], "trusted ca file")
-        self.assertEqual(kwargs["_tls_client_cert_key_file"], "trusted client cert")
-        self.assertEqual(kwargs["_tls_client_cert_key_password"], "key password")
+        assert kwargs["_tls_verify_hostname"] == "hostname"
+        assert kwargs["_tls_trusted_ca_file"] == "trusted ca file"
+        assert kwargs["_tls_client_cert_key_file"] == "trusted client cert"
+        assert kwargs["_tls_client_cert_key_password"] == "key password"
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_useragent_header(self, mock_client_class):
@@ -100,7 +100,7 @@ def test_useragent_header(self, mock_client_class):
             "User-Agent",
             "{}/{}".format(databricks.sql.USER_AGENT_NAME, databricks.sql.__version__),
         )
-        self.assertIn(user_agent_header, http_headers)
+        assert user_agent_header in http_headers
 
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, user_agent_entry="foobar")
         user_agent_header_with_entry = (
@@ -110,7 +110,7 @@ def test_useragent_header(self, mock_client_class):
             ),
         )
         http_headers = mock_client_class.call_args[0][3]
-        self.assertIn(user_agent_header_with_entry, http_headers)
+        assert user_agent_header_with_entry in http_headers
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_context_manager_closes_connection(self, mock_client_class):
@@ -125,13 +125,13 @@ def test_context_manager_closes_connection(self, mock_client_class):
 
         # Check that close_session was called with the correct SessionId
         close_session_call_args = instance.close_session.call_args[0][0]
-        self.assertEqual(close_session_call_args.guid, b"\x22")
-        self.assertEqual(close_session_call_args.secret, b"\x33")
+        assert close_session_call_args.guid == b"\x22"
+        assert close_session_call_args.secret == b"\x33"
 
         connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         connection.close = Mock()
         try:
-            with self.assertRaises(KeyboardInterrupt):
+            with pytest.raises(KeyboardInterrupt):
                 with connection:
                     raise KeyboardInterrupt("Simulated interrupt")
         finally:
@@ -143,14 +143,12 @@ def test_max_number_of_retries_passthrough(self, mock_client_class):
             _retry_stop_after_attempts_count=54, **self.DUMMY_CONNECTION_ARGS
         )
 
-        self.assertEqual(
-            mock_client_class.call_args[1]["_retry_stop_after_attempts_count"], 54
-        )
+        assert mock_client_class.call_args[1]["_retry_stop_after_attempts_count"] == 54
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_socket_timeout_passthrough(self, mock_client_class):
         databricks.sql.connect(_socket_timeout=234, **self.DUMMY_CONNECTION_ARGS)
-        self.assertEqual(mock_client_class.call_args[1]["_socket_timeout"], 234)
+        assert mock_client_class.call_args[1]["_socket_timeout"] == 234
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
@@ -160,7 +158,7 @@ def test_configuration_passthrough(self, mock_client_class):
         )
 
         call_kwargs = mock_client_class.return_value.open_session.call_args[1]
-        self.assertEqual(call_kwargs["session_configuration"], mock_session_config)
+        assert call_kwargs["session_configuration"] == mock_session_config
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
@@ -171,8 +169,8 @@ def test_initial_namespace_passthrough(self, mock_client_class):
         )
 
         call_kwargs = mock_client_class.return_value.open_session.call_args[1]
-        self.assertEqual(call_kwargs["catalog"], mock_cat)
-        self.assertEqual(call_kwargs["schema"], mock_schem)
+        assert call_kwargs["catalog"] == mock_cat
+        assert call_kwargs["schema"] == mock_schem
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):
@@ -188,9 +186,5 @@ def test_finalizer_closes_abandoned_connection(self, mock_client_class):
 
         # Check that close_session was called with the correct SessionId
         close_session_call_args = instance.close_session.call_args[0][0]
-        self.assertEqual(close_session_call_args.guid, b"\x22")
-        self.assertEqual(close_session_call_args.secret, b"\x33")
-
-
-if __name__ == "__main__":
-    unittest.main()
+        assert close_session_call_args.guid == b"\x22"
+        assert close_session_call_args.secret == b"\x33"