convert test to pytest

Signed-off-by: Vaibhav <vrongmeal@gmail.com>

justfile

@@ -123,10 +123,12 @@ venv:
   {{VENV_BIN}}/python -m pip install --upgrade pip
   {{VENV_BIN}}/python -m pip install poetry
 
-# Runs pytest in the tests directory.
-pytest *args:
-  {{VENV_BIN}}/poetry -C tests install --no-root
+poetry:
   {{VENV_BIN}}/poetry -C tests lock --no-update
+  {{VENV_BIN}}/poetry -C tests install --no-root
+
+# Runs pytest in the tests directory.
+pytest *args: poetry
   {{VENV_BIN}}/poetry -C tests run pytest -v --rootdir={{invocation_directory()}}/tests {{ if args == "" {'tests'} else {args} }}
 
 # private helpers below

tests/conftest.py

@@ -1 +1,4 @@
-pytest_plugins = ["tests.fixtures.glaredb"]
+pytest_plugins = [
+    "tests.fixtures.glaredb",
+    "tests.fixtures.iceberg",
+]

tests/fixtures/glaredb.py

@@ -11,7 +11,7 @@
 
 import tests
 
-logger = logging.getLogger("fixtures")
+logger = logging.getLogger("fixtures.glaredb")
 
 
 @pytest.fixture

tests/fixtures/iceberg.py

@@ -0,0 +1,60 @@
+import logging
+import pathlib
+
+import pytest
+
+from pyiceberg.catalog.sql import SqlCatalog
+from pyiceberg.exceptions import NamespaceAlreadyExistsError, TableAlreadyExistsError
+
+import pyarrow.parquet as pq
+
+import tests
+
+
+logger = logging.getLogger("fixtures.iceberg")
+
+
+@pytest.fixture
+def pyiceberg_table() -> pathlib.Path:
+    warehouse_path = tests.PKG_DIRECTORY.joinpath("testdata", "iceberg").absolute()
+    taxi_data_source_path = warehouse_path.joinpath(
+        "source_data",
+        "yellow_tripdata_2023-01.parquet",
+    )
+
+    catalog = SqlCatalog(
+        "default",
+        **{
+            "uri": f"sqlite:///{warehouse_path}/pyiceberg_catalog.db",
+            "warehouse": f"file://{warehouse_path}",
+            "py-io-impl": "pyiceberg.io.fsspec.FsspecFileIO"
+        },
+    )
+
+    df = pq.read_table(taxi_data_source_path)
+
+    NAMESPACE = "pyiceberg"
+    db_dir = warehouse_path.joinpath(f"{NAMESPACE}.db")
+
+    try:
+        catalog.create_namespace(NAMESPACE)
+    except NamespaceAlreadyExistsError:
+        logger.warning(f"Catalog '{NAMESPACE}' already exists")
+
+    TABLE_NAME = "taxi_dataset"
+
+    try:
+        table = catalog.create_table(
+            f"{NAMESPACE}.{TABLE_NAME}",
+            schema=df.schema,
+        )
+        table.overwrite(df)
+        num_rows = len(table.scan().to_arrow())
+
+        logger.info(f"Table '{TABLE_NAME}' created with {num_rows} rows.")
+
+    except TableAlreadyExistsError:
+        logger.warning(f"Table '{TABLE_NAME}' already exists in the catalog.")
+        logger.info(f"To re-create, remove the existing table at {db_dir}")
+
+    return db_dir.joinpath(TABLE_NAME)

tests/pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 python = "^3.11"
 pytest = "^7.4.3"
 pymongo = "^4.6.3"
-psycopg2 = "^2.9.9"
+psycopg2-binary = "^2.9.9"
 pyarrow = "^14.0.2"
 pylance = "^0.9.6"
 ruff = "0.1.14"
@@ -18,6 +18,8 @@ dbt-postgres = "^1.7.7"
 pytest-timeout = "^2.2.0"
 pytest-xdist = "^3.5.0"
 pytest-benchmark = "^4.0.0"
+botocore = "^1.34.81"
+pyiceberg = { version = "^0.6.0", extras = ["pyarrow", "sql-sqlite"] }
 
 [build-system]
 requires = ["poetry-core"]

tests/test_iceberg.py

@@ -0,0 +1,14 @@
+import pathlib
+
+import psycopg2
+
+
+def test_pyiceberg(
+    pyiceberg_table: pathlib.Path,
+    glaredb_connection: psycopg2.extensions.connection,
+):
+    with glaredb_connection.cursor() as cur:
+        cur.execute(f"select count(*) from read_iceberg('{pyiceberg_table}')")
+        result = cur.fetchall()
+
+    print(result)