6
6
"""
7
7
8
8
import pytest
9
- from unittest .mock import patch , MagicMock , Mock
9
+ from unittest .mock import Mock
10
10
11
11
from databricks .sql .result_set import SeaResultSet , Row
12
12
from databricks .sql .utils import JsonQueue
13
- from databricks .sql .backend .types import CommandId , CommandState , BackendType
13
+ from databricks .sql .backend .types import CommandId , CommandState
14
14
from databricks .sql .backend .sea .models .base import ResultData , ResultManifest
15
15
16
16
@@ -88,6 +88,18 @@ def json_queue(self, sample_data):
88
88
"""Create a JsonQueue with sample data."""
89
89
return JsonQueue (sample_data )
90
90
91
+ def empty_manifest (self ):
92
+ """Create an empty manifest."""
93
+ return ResultManifest (
94
+ format = "JSON_ARRAY" ,
95
+ schema = {},
96
+ total_row_count = 0 ,
97
+ total_byte_count = 0 ,
98
+ total_chunk_count = 0 ,
99
+ truncated = False ,
100
+ is_volume_operation = False ,
101
+ )
102
+
91
103
def test_init_with_execute_response (
92
104
self , mock_connection , mock_sea_client , execute_response
93
105
):
@@ -97,6 +109,7 @@ def test_init_with_execute_response(
97
109
execute_response = execute_response ,
98
110
sea_client = mock_sea_client ,
99
111
result_data = ResultData (data = []),
112
+ manifest = self .empty_manifest (),
100
113
buffer_size_bytes = 1000 ,
101
114
arraysize = 100 ,
102
115
)
@@ -117,6 +130,7 @@ def test_close(self, mock_connection, mock_sea_client, execute_response):
117
130
execute_response = execute_response ,
118
131
sea_client = mock_sea_client ,
119
132
result_data = ResultData (data = []),
133
+ manifest = self .empty_manifest (),
120
134
buffer_size_bytes = 1000 ,
121
135
arraysize = 100 ,
122
136
)
@@ -138,6 +152,7 @@ def test_close_when_already_closed_server_side(
138
152
execute_response = execute_response ,
139
153
sea_client = mock_sea_client ,
140
154
result_data = ResultData (data = []),
155
+ manifest = self .empty_manifest (),
141
156
buffer_size_bytes = 1000 ,
142
157
arraysize = 100 ,
143
158
)
@@ -161,6 +176,7 @@ def test_close_when_connection_closed(
161
176
execute_response = execute_response ,
162
177
sea_client = mock_sea_client ,
163
178
result_data = ResultData (data = []),
179
+ manifest = self .empty_manifest (),
164
180
buffer_size_bytes = 1000 ,
165
181
arraysize = 100 ,
166
182
)
@@ -317,6 +333,7 @@ def test_fetchmany_arrow_not_implemented(
317
333
execute_response = execute_response ,
318
334
sea_client = mock_sea_client ,
319
335
result_data = ResultData (data = None , external_links = []),
336
+ manifest = self .empty_manifest (),
320
337
buffer_size_bytes = 1000 ,
321
338
arraysize = 100 ,
322
339
)
@@ -336,6 +353,7 @@ def test_fetchall_arrow_not_implemented(
336
353
execute_response = execute_response ,
337
354
sea_client = mock_sea_client ,
338
355
result_data = ResultData (data = None , external_links = []),
356
+ manifest = self .empty_manifest (),
339
357
buffer_size_bytes = 1000 ,
340
358
arraysize = 100 ,
341
359
)
@@ -353,6 +371,7 @@ def test_is_staging_operation(
353
371
execute_response = execute_response ,
354
372
sea_client = mock_sea_client ,
355
373
result_data = ResultData (data = []),
374
+ manifest = self .empty_manifest (),
356
375
buffer_size_bytes = 1000 ,
357
376
arraysize = 100 ,
358
377
)
0 commit comments