fix int test types

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -867,9 +867,7 @@ def get_execution_result(
             ssl_options=self._ssl_options,
         )
 
-        status = CommandState.from_thrift_state(resp.status)
-        if status is None:
-            raise ValueError(f"Invalid operation state: {resp.status}")
+        status = self.get_query_state(command_id)
 
         execute_response = ExecuteResponse(
             command_id=command_id,

tests/e2e/common/retry_test_mixins.py

@@ -326,7 +326,7 @@ def test_retry_abort_close_operation_on_404(self, caplog):
         with self.connection(extra_params={**self._retry_policy}) as conn:
             with conn.cursor() as curs:
                 with patch(
-                    "databricks.sql.utils.ExecuteResponse.has_been_closed_server_side",
+                    "databricks.sql.backend.types.ExecuteResponse.has_been_closed_server_side",
                     new_callable=PropertyMock,
                     return_value=False,
                 ):

tests/e2e/test_driver.py

@@ -933,12 +933,12 @@ def test_result_set_close(self):
                 result_set = cursor.active_result_set
                 assert result_set is not None
 
-                initial_op_state = result_set.op_state
+                initial_op_state = result_set.status
 
                 result_set.close()
 
-                assert result_set.op_state == CommandState.CLOSED
-                assert result_set.op_state != initial_op_state
+                assert result_set.status == CommandState.CLOSED
+                assert result_set.status != initial_op_state
 
                 # Closing the result set again should be a no-op and not raise exceptions
                 result_set.close()