forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_column_queue.py
More file actions
26 lines (20 loc) · 974 Bytes
/
test_column_queue.py
File metadata and controls
26 lines (20 loc) · 974 Bytes
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
from databricks.sql.utils import ColumnQueue, ColumnTable
class TestColumnQueueSuite:
@staticmethod
def make_column_table(table):
n_cols = len(table) if table else 0
return ColumnTable(table, [f"col_{i}" for i in range(n_cols)])
def test_fetchmany_respects_n_rows(self):
column_table = self.make_column_table(
[[0, 3, 6, 9], [1, 4, 7, 10], [2, 5, 8, 11]]
)
column_queue = ColumnQueue(column_table)
assert column_queue.next_n_rows(2) == column_table.slice(0, 2)
assert column_queue.next_n_rows(2) == column_table.slice(2, 2)
def test_fetch_remaining_rows_respects_n_rows(self):
column_table = self.make_column_table(
[[0, 3, 6, 9], [1, 4, 7, 10], [2, 5, 8, 11]]
)
column_queue = ColumnQueue(column_table)
assert column_queue.next_n_rows(2) == column_table.slice(0, 2)
assert column_queue.remaining_rows() == column_table.slice(2, 2)