-
Notifications
You must be signed in to change notification settings - Fork 140
Expand file tree
/
Copy pathtest_filters.py
More file actions
157 lines (132 loc) · 6.33 KB
/
test_filters.py
File metadata and controls
157 lines (132 loc) · 6.33 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
"""
Tests for the ResultSetFilter class.
"""
import unittest
from unittest.mock import MagicMock, patch
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
class TestResultSetFilter(unittest.TestCase):
"""Tests for the ResultSetFilter class."""
def setUp(self):
"""Set up test fixtures."""
# Create a mock SeaResultSet
self.mock_sea_result_set = MagicMock()
# Set up the remaining_rows method on the results attribute
self.mock_sea_result_set.results = MagicMock()
self.mock_sea_result_set.results.remaining_rows.return_value = [
["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""],
["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""],
[
"catalog1",
"schema1",
"table3",
"owner1",
"2023-01-01",
"SYSTEM TABLE",
"",
],
[
"catalog1",
"schema1",
"table4",
"owner1",
"2023-01-01",
"EXTERNAL TABLE",
"",
],
]
# Set up the connection and other required attributes
self.mock_sea_result_set.connection = MagicMock()
self.mock_sea_result_set.backend = MagicMock()
self.mock_sea_result_set.buffer_size_bytes = 1000
self.mock_sea_result_set.arraysize = 100
self.mock_sea_result_set.statement_id = "test-statement-id"
self.mock_sea_result_set.lz4_compressed = False
# Create a mock CommandId
from databricks.sql.backend.types import CommandId, BackendType
mock_command_id = CommandId(BackendType.SEA, "test-statement-id")
self.mock_sea_result_set.command_id = mock_command_id
self.mock_sea_result_set.status = MagicMock()
self.mock_sea_result_set.description = [
("catalog_name", "string", None, None, None, None, True),
("schema_name", "string", None, None, None, None, True),
("table_name", "string", None, None, None, None, True),
("owner", "string", None, None, None, None, True),
("creation_time", "string", None, None, None, None, True),
("table_type", "string", None, None, None, None, True),
("remarks", "string", None, None, None, None, True),
]
self.mock_sea_result_set.has_been_closed_server_side = False
self.mock_sea_result_set._arrow_schema_bytes = None
def test__filter_json_result_set(self):
"""Test filtering by column values with various options."""
# Case 1: Case-sensitive filtering
allowed_values = ["table1", "table3"]
with patch(
"databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
):
with patch(
"databricks.sql.backend.sea.result_set.SeaResultSet"
) as mock_sea_result_set_class:
mock_instance = MagicMock()
mock_sea_result_set_class.return_value = mock_instance
# Call _filter_json_result_set on the table_name column (index 2)
result = ResultSetFilter._filter_json_result_set(
self.mock_sea_result_set, 2, allowed_values, case_sensitive=True
)
# Verify the filter was applied correctly
mock_sea_result_set_class.assert_called_once()
# Check the filtered data passed to the constructor
args, kwargs = mock_sea_result_set_class.call_args
result_data = kwargs.get("result_data")
self.assertIsNotNone(result_data)
self.assertEqual(len(result_data.data), 2)
self.assertIn(result_data.data[0][2], allowed_values)
self.assertIn(result_data.data[1][2], allowed_values)
# Case 2: Case-insensitive filtering
mock_sea_result_set_class.reset_mock()
with patch(
"databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
):
with patch(
"databricks.sql.backend.sea.result_set.SeaResultSet"
) as mock_sea_result_set_class:
mock_instance = MagicMock()
mock_sea_result_set_class.return_value = mock_instance
# Call _filter_json_result_set with case-insensitive matching
result = ResultSetFilter._filter_json_result_set(
self.mock_sea_result_set,
2,
["TABLE1", "TABLE3"],
case_sensitive=False,
)
mock_sea_result_set_class.assert_called_once()
def test_filter_tables_by_type(self):
"""Test filtering tables by type with various options."""
# Case 1: Specific table types
table_types = ["TABLE", "VIEW"]
# Mock results as JsonQueue (not CloudFetchQueue or ArrowQueue)
from databricks.sql.backend.sea.queue import JsonQueue
self.mock_sea_result_set.results = JsonQueue([])
with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter:
ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types)
args, kwargs = mock_filter.call_args
self.assertEqual(args[0], self.mock_sea_result_set)
self.assertEqual(kwargs.get("column_index"), 5) # Table type column index
self.assertEqual(kwargs.get("allowed_values"), table_types)
self.assertEqual(kwargs.get("case_sensitive"), True)
# Case 2: Default table types (None or empty list)
with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter:
# Test with None
ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
args, kwargs = mock_filter.call_args
self.assertEqual(
kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"]
)
# Test with empty list
ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
args, kwargs = mock_filter.call_args
self.assertEqual(
kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"]
)
if __name__ == "__main__":
unittest.main()