Added e2e test

nithinkdb · nithinkdb · commit 031c4e87799a · 2023-10-16T11:55:51.000-07:00
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
@@ -225,8 +225,6 @@ def __del__(self):
     def get_session_id(self):
         return self.thrift_backend.handle_to_id(self._session_handle)
 
-<<<<<<< HEAD
-=======
     def get_session_protocol_version(self):
         return self.thrift_backend.extract_protocol_version_from_handle(
             self._session_handle
@@ -239,7 +237,6 @@ def server_parameterized_queries_enabled(protocolVersion):
         else:
             return False
 
->>>>>>> cd93869 (Put in some unit tests, will add e2e)
     def get_session_id_hex(self):
         return self.thrift_backend.handle_to_hex_id(self._session_handle)
 
@@ -516,13 +513,11 @@ def execute(
         """
         if parameters is None:
             parameters = []
-<<<<<<< HEAD
-=======
-        elif (not self.server_parameterized_queries_enabled(self.connection._session_handle)):
+
+        elif (not Connection.server_parameterized_queries_enabled(self.connection._session_handle)):
             raise Error(
-                "Parameterized operations are not supported by this server. DBR 14.2 is required."
+                "Parameterized operations are not supported by this server. DBR 14.1 is required."
             )
->>>>>>> cd93869 (Put in some unit tests, will add e2e)
         else:
             parameters = named_parameters_to_tsparkparams(parameters)
 
diff --git a/tests/e2e/common/parameterized_query_tests.py b/tests/e2e/common/parameterized_query_tests.py
@@ -2,8 +2,10 @@
 from decimal import Decimal
 from enum import Enum
 from typing import Dict, List, Tuple, Union
+from unittest.mock import Mock, patch, MagicMock
+from databricks.sql import Error
 
-import pytz
+import pytest
 
 from databricks.sql.exc import DatabaseError
 from databricks.sql.utils import (
@@ -29,19 +31,18 @@ def _get_one_result(self, query: str, parameters: Union[Dict, List[Dict]]) -> Tu
             with conn.cursor() as cursor:
                 cursor.execute(query, parameters=parameters)
                 return cursor.fetchone()
-            
-    def _test_protocol_checking(self) -> Tuple:
-        with self.connection() as conn:
-            with conn.cursor() as cursor:
-                cursor.execute(query, parameters=parameters)
-                return cursor.fetchone()
-
+        
     def _quantize(self, input: Union[float, int], place_value=2) -> Decimal:
 
         return Decimal(str(input)).quantize(Decimal("0." + "0" * place_value))
 
+    @patch('databricks.sql.client.Connection.server_parameterized_queries_enabled', return_value=False)
+    def test_protocol_too_low(self, mock_parameterized_queries_enabled):
+        params = {"p": None}
+        with pytest.raises(Error, match="Parameterized operations are not supported by this server. DBR 14.1 is required."):
+            result = self._get_one_result(self.QUERY, params)
+    
     def test_primitive_inferred_none(self):
-
         params = {"p": None}
         result = self._get_one_result(self.QUERY, params)
         assert result.col == None
