Skip to content

Commit b558bc8

Browse files
committed
unit test fix
Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
1 parent 27295c2 commit b558bc8

File tree

1 file changed

+16
-12
lines changed

1 file changed

+16
-12
lines changed

tests/unit/test_fetches.py

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,25 @@
11
import unittest
22
import pytest
3-
from unittest.mock import Mock
3+
from unittest.mock import Mock, patch
44

55
try:
66
import pyarrow as pa
77
except ImportError:
88
pa = None
99

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+
1023
import databricks.sql.client as client
1124
from databricks.sql.utils import ExecuteResponse, ArrowQueue
1225

@@ -31,15 +44,6 @@ def make_arrow_queue(batch):
3144
_, table = FetchTests.make_arrow_table(batch)
3245
queue = ArrowQueue(table, len(batch))
3346
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
4347

4448
@staticmethod
4549
def make_dummy_result_set_from_initial_results(initial_results):
@@ -48,7 +52,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
4852
arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
4953
rs = client.ResultSet(
5054
connection=Mock(),
51-
thrift_backend=FetchTests.mock_thrift_backend_with_retry_policy(),
55+
thrift_backend=None,
5256
execute_response=ExecuteResponse(
5357
status=None,
5458
has_been_closed_server_side=True,
@@ -88,7 +92,7 @@ def fetch_results(
8892

8993
return results, batch_index < len(batch_list)
9094

91-
mock_thrift_backend = FetchTests.mock_thrift_backend_with_retry_policy()
95+
mock_thrift_backend = Mock()
9296
mock_thrift_backend.fetch_results = fetch_results
9397
num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
9498

0 commit comments

Comments
 (0)