-
Notifications
You must be signed in to change notification settings - Fork 95
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added the unit test for column_queue
- Loading branch information
1 parent
0ddca9d
commit 4351410
Showing
1 changed file
with
24 additions
and
0 deletions.
There are no files selected for viewing
24 changes: 24 additions & 0 deletions
24
databricks_sql_connector_core/tests/unit/test_column_queue.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
import unittest | ||
import pytest | ||
import pyarrow as pa | ||
from databricks_sql_connector_core.sql.utils import ColumnQueue | ||
|
||
|
||
class TestColumnQueueSuite: | ||
@pytest.fixture(scope="function") | ||
def setup(self): | ||
columnar_table = [[0, 3, 6, 9], [1, 4, 7, 10], [2, 5, 8, 11]] | ||
column_names = [f"col_{i}" for i in range(len(columnar_table))] | ||
return ColumnQueue(columnar_table, column_names) | ||
|
||
def test_fetchmany_respects_n_rows(self, setup): | ||
column_queue = setup | ||
assert column_queue.next_n_rows(2) == [[0, 3], [1, 4], [2, 5]] | ||
assert column_queue.next_n_rows(2) == [[6, 9], [7, 10], [8, 11]] | ||
|
||
def test_fetch_remaining_rows_respects_n_rows(self, setup): | ||
column_queue = setup | ||
assert column_queue.next_n_rows(2) == [[0, 3], [1, 4], [2, 5]] | ||
assert column_queue.remaining_rows() == [[6, 9], [7, 10], [8, 11]] | ||
|
||
|