From 786dc1e7ce169169882cdfb9923e0c22b4c81a3c Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Wed, 14 Aug 2024 14:49:41 +0530
Subject: [PATCH] Added the unit test for column_queue

Fixed __version__

Fix
---
 .../conftest.py                               |  4 ++++
 .../tests/e2e/common/predicates.py            |  4 ++--
 .../tests/unit/test_column_queue.py           | 24 +++++++++++++++++++
 setup_script.py                               |  1 +
 4 files changed, 31 insertions(+), 2 deletions(-)
 rename conftest.py => databricks_sql_connector_core/conftest.py (86%)
 create mode 100644 databricks_sql_connector_core/tests/unit/test_column_queue.py

diff --git a/conftest.py b/databricks_sql_connector_core/conftest.py
similarity index 86%
rename from conftest.py
rename to databricks_sql_connector_core/conftest.py
index c8b350be..3f69753f 100644
--- a/conftest.py
+++ b/databricks_sql_connector_core/conftest.py
@@ -4,11 +4,13 @@
 
 @pytest.fixture(scope="session")
 def host():
+    return "adb-6436897454825492.12.azuredatabricks.net"
     return os.getenv("DATABRICKS_SERVER_HOSTNAME")
 
 
 @pytest.fixture(scope="session")
 def http_path():
+    return "/sql/1.0/warehouses/2f03dd43e35e2aa0"
     return os.getenv("DATABRICKS_HTTP_PATH")
 
 
@@ -24,11 +26,13 @@ def ingestion_user():
 
 @pytest.fixture(scope="session")
 def catalog():
+    return "main"
     return os.getenv("DATABRICKS_CATALOG")
 
 
 @pytest.fixture(scope="session")
 def schema():
+    return "default"
     return os.getenv("DATABRICKS_SCHEMA", "default")
 
 
diff --git a/databricks_sql_connector_core/tests/e2e/common/predicates.py b/databricks_sql_connector_core/tests/e2e/common/predicates.py
index 10dbde9e..529fb7fe 100644
--- a/databricks_sql_connector_core/tests/e2e/common/predicates.py
+++ b/databricks_sql_connector_core/tests/e2e/common/predicates.py
@@ -18,13 +18,13 @@ def pysql_has_version(compare, version):
 
     Expected use:
         from common.predicates import pysql_has_version
-        from databricks import sql as pysql
+        from databricks_sql_connector_core import sql as pysql
         ...
         @unittest.skipIf(pysql_has_version('<', '2'))
         def test_some_pyhive_v1_stuff():
             ...
     """
-    from databricks import sql
+    from databricks_sql_connector_core import sql
     return compare_module_version(sql, compare, version)
 
 
diff --git a/databricks_sql_connector_core/tests/unit/test_column_queue.py b/databricks_sql_connector_core/tests/unit/test_column_queue.py
new file mode 100644
index 00000000..dd780639
--- /dev/null
+++ b/databricks_sql_connector_core/tests/unit/test_column_queue.py
@@ -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]]
+
+
diff --git a/setup_script.py b/setup_script.py
index 14929888..27b021cb 100644
--- a/setup_script.py
+++ b/setup_script.py
@@ -7,6 +7,7 @@ def build_and_install_library(directory_name):
     # Change directory to one level down
     os.chdir(directory_name)
 
+    # chal jaa
     # Build the library using Poetry
     subprocess.run(['poetry', 'build'], check=True)