1
1
import unittest
2
2
import pytest
3
- from unittest .mock import Mock
3
+ from unittest .mock import Mock , patch
4
4
5
5
try :
6
6
import pyarrow as pa
7
7
except ImportError :
8
8
pa = None
9
9
10
+ def noop_log_latency_decorator (* args , ** kwargs ):
11
+ """
12
+ This is a no-op decorator. It is used to patch the log_latency decorator
13
+ during tests, so that the tests for client logic are not affected by the
14
+ telemetry logging logic. It accepts any arguments and returns a decorator
15
+ that returns the original function unmodified.
16
+ """
17
+ def decorator (func ):
18
+ return func
19
+ return decorator
20
+
21
+ patch ('databricks.sql.telemetry.latency_logger.log_latency' , new = noop_log_latency_decorator ).start ()
22
+
10
23
import databricks .sql .client as client
11
24
from databricks .sql .utils import ExecuteResponse , ArrowQueue
12
25
@@ -31,15 +44,6 @@ def make_arrow_queue(batch):
31
44
_ , table = FetchTests .make_arrow_table (batch )
32
45
queue = ArrowQueue (table , len (batch ))
33
46
return queue
34
-
35
- @classmethod
36
- def mock_thrift_backend_with_retry_policy (cls ): # Required for log_latency() decorator
37
- """Create a simple thrift_backend mock with retry_policy for basic tests."""
38
- mock_thrift_backend = Mock ()
39
- mock_retry_policy = Mock ()
40
- mock_retry_policy .history = []
41
- mock_thrift_backend .retry_policy = mock_retry_policy
42
- return mock_thrift_backend
43
47
44
48
@staticmethod
45
49
def make_dummy_result_set_from_initial_results (initial_results ):
@@ -48,7 +52,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
48
52
arrow_queue = ArrowQueue (arrow_table , len (initial_results ), 0 )
49
53
rs = client .ResultSet (
50
54
connection = Mock (),
51
- thrift_backend = FetchTests . mock_thrift_backend_with_retry_policy () ,
55
+ thrift_backend = None ,
52
56
execute_response = ExecuteResponse (
53
57
status = None ,
54
58
has_been_closed_server_side = True ,
@@ -88,7 +92,7 @@ def fetch_results(
88
92
89
93
return results , batch_index < len (batch_list )
90
94
91
- mock_thrift_backend = FetchTests . mock_thrift_backend_with_retry_policy ()
95
+ mock_thrift_backend = Mock ()
92
96
mock_thrift_backend .fetch_results = fetch_results
93
97
num_cols = len (batch_list [0 ][0 ]) if batch_list and batch_list [0 ] else 0
94
98
0 commit comments