Skip to content

Commit 28675f5

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
introduce unit tests for metadata methods
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent f6d873d commit 28675f5

File tree

2 files changed

+342
-44
lines changed

2 files changed

+342
-44
lines changed

‎tests/unit/test_filters.py

Lines changed: 90 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -15,24 +15,39 @@ def setUp(self):
1515
"""Set up test fixtures."""
1616
# Create a mock SeaResultSet
1717
self.mock_sea_result_set = MagicMock()
18-
self.mock_sea_result_set._response = {
19-
"result": {
20-
"data_array": [
21-
["catalog1", "schema1", "table1", "TABLE", ""],
22-
["catalog1", "schema1", "table2", "VIEW", ""],
23-
["catalog1", "schema1", "table3", "SYSTEM TABLE", ""],
24-
["catalog1", "schema1", "table4", "EXTERNAL TABLE", ""],
25-
],
26-
"row_count": 4,
27-
}
28-
}
18+
19+
# Set up the remaining_rows method on the results attribute
20+
self.mock_sea_result_set.results = MagicMock()
21+
self.mock_sea_result_set.results.remaining_rows.return_value = [
22+
["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""],
23+
["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""],
24+
[
25+
"catalog1",
26+
"schema1",
27+
"table3",
28+
"owner1",
29+
"2023-01-01",
30+
"SYSTEM TABLE",
31+
"",
32+
],
33+
[
34+
"catalog1",
35+
"schema1",
36+
"table4",
37+
"owner1",
38+
"2023-01-01",
39+
"EXTERNAL TABLE",
40+
"",
41+
],
42+
]
2943

3044
# Set up the connection and other required attributes
3145
self.mock_sea_result_set.connection = MagicMock()
3246
self.mock_sea_result_set.backend = MagicMock()
3347
self.mock_sea_result_set.buffer_size_bytes = 1000
3448
self.mock_sea_result_set.arraysize = 100
3549
self.mock_sea_result_set.statement_id = "test-statement-id"
50+
self.mock_sea_result_set.lz4_compressed = False
3651

3752
# Create a mock CommandId
3853
from databricks.sql.backend.types import CommandId, BackendType
@@ -45,70 +60,102 @@ def setUp(self):
4560
("catalog_name", "string", None, None, None, None, True),
4661
("schema_name", "string", None, None, None, None, True),
4762
("table_name", "string", None, None, None, None, True),
63+
("owner", "string", None, None, None, None, True),
64+
("creation_time", "string", None, None, None, None, True),
4865
("table_type", "string", None, None, None, None, True),
4966
("remarks", "string", None, None, None, None, True),
5067
]
5168
self.mock_sea_result_set.has_been_closed_server_side = False
69+
self.mock_sea_result_set._arrow_schema_bytes = None
5270

53-
def test_filter_tables_by_type(self):
54-
"""Test filtering tables by type."""
55-
# Test with specific table types
56-
table_types = ["TABLE", "VIEW"]
71+
def test_filter_by_column_values(self):
72+
"""Test filtering by column values with various options."""
73+
# Case 1: Case-sensitive filtering
74+
allowed_values = ["table1", "table3"]
5775

58-
# Make the mock_sea_result_set appear to be a SeaResultSet
5976
with patch("databricks.sql.backend.filters.isinstance", return_value=True):
6077
with patch(
6178
"databricks.sql.result_set.SeaResultSet"
6279
) as mock_sea_result_set_class:
63-
# Set up the mock to return a new mock when instantiated
6480
mock_instance = MagicMock()
6581
mock_sea_result_set_class.return_value = mock_instance
6682

67-
result = ResultSetFilter.filter_tables_by_type(
68-
self.mock_sea_result_set, table_types
83+
# Call filter_by_column_values on the table_name column (index 2)
84+
result = ResultSetFilter.filter_by_column_values(
85+
self.mock_sea_result_set, 2, allowed_values, case_sensitive=True
6986
)
7087

7188
# Verify the filter was applied correctly
7289
mock_sea_result_set_class.assert_called_once()
7390

74-
def test_filter_tables_by_type_case_insensitive(self):
75-
"""Test filtering tables by type with case insensitivity."""
76-
# Test with lowercase table types
77-
table_types = ["table", "view"]
91+
# Check the filtered data passed to the constructor
92+
args, kwargs = mock_sea_result_set_class.call_args
93+
result_data = kwargs.get("result_data")
94+
self.assertIsNotNone(result_data)
95+
self.assertEqual(len(result_data.data), 2)
96+
self.assertIn(result_data.data[0][2], allowed_values)
97+
self.assertIn(result_data.data[1][2], allowed_values)
7898

79-
# Make the mock_sea_result_set appear to be a SeaResultSet
99+
# Case 2: Case-insensitive filtering
100+
mock_sea_result_set_class.reset_mock()
80101
with patch("databricks.sql.backend.filters.isinstance", return_value=True):
81102
with patch(
82103
"databricks.sql.result_set.SeaResultSet"
83104
) as mock_sea_result_set_class:
84-
# Set up the mock to return a new mock when instantiated
85105
mock_instance = MagicMock()
86106
mock_sea_result_set_class.return_value = mock_instance
87107

88-
result = ResultSetFilter.filter_tables_by_type(
89-
self.mock_sea_result_set, table_types
108+
# Call filter_by_column_values with case-insensitive matching
109+
result = ResultSetFilter.filter_by_column_values(
110+
self.mock_sea_result_set,
111+
2,
112+
["TABLE1", "TABLE3"],
113+
case_sensitive=False,
90114
)
91-
92-
# Verify the filter was applied correctly
93115
mock_sea_result_set_class.assert_called_once()
94116

95-
def test_filter_tables_by_type_default(self):
96-
"""Test filtering tables by type with default types."""
97-
# Make the mock_sea_result_set appear to be a SeaResultSet
98-
with patch("databricks.sql.backend.filters.isinstance", return_value=True):
99-
with patch(
100-
"databricks.sql.result_set.SeaResultSet"
101-
) as mock_sea_result_set_class:
102-
# Set up the mock to return a new mock when instantiated
103-
mock_instance = MagicMock()
104-
mock_sea_result_set_class.return_value = mock_instance
117+
# Case 3: Unsupported result set type
118+
mock_unsupported_result_set = MagicMock()
119+
with patch("databricks.sql.backend.filters.isinstance", return_value=False):
120+
with patch("databricks.sql.backend.filters.logger") as mock_logger:
121+
result = ResultSetFilter.filter_by_column_values(
122+
mock_unsupported_result_set, 0, ["value"], True
123+
)
124+
mock_logger.warning.assert_called_once()
125+
self.assertEqual(result, mock_unsupported_result_set)
126+
127+
def test_filter_tables_by_type(self):
128+
"""Test filtering tables by type with various options."""
129+
# Case 1: Specific table types
130+
table_types = ["TABLE", "VIEW"]
105131

106-
result = ResultSetFilter.filter_tables_by_type(
107-
self.mock_sea_result_set, None
132+
with patch("databricks.sql.backend.filters.isinstance", return_value=True):
133+
with patch.object(
134+
ResultSetFilter, "filter_by_column_values"
135+
) as mock_filter:
136+
ResultSetFilter.filter_tables_by_type(
137+
self.mock_sea_result_set, table_types
108138
)
139+
args, kwargs = mock_filter.call_args
140+
self.assertEqual(args[0], self.mock_sea_result_set)
141+
self.assertEqual(args[1], 5) # Table type column index
142+
self.assertEqual(args[2], table_types)
143+
self.assertEqual(kwargs.get("case_sensitive"), True)
109144

110-
# Verify the filter was applied correctly
111-
mock_sea_result_set_class.assert_called_once()
145+
# Case 2: Default table types (None or empty list)
146+
with patch("databricks.sql.backend.filters.isinstance", return_value=True):
147+
with patch.object(
148+
ResultSetFilter, "filter_by_column_values"
149+
) as mock_filter:
150+
# Test with None
151+
ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
152+
args, kwargs = mock_filter.call_args
153+
self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
154+
155+
# Test with empty list
156+
ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
157+
args, kwargs = mock_filter.call_args
158+
self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
112159

113160

114161
if __name__ == "__main__":

0 commit comments

Comments
 (0)