Using fixture for catalog

Author: datapythonista

pandas/tests/io/test_iceberg.py

@@ -5,10 +5,9 @@
 data used for Parquet tests (``pandas/tests/io/data/parquet/simple.parquet``).
 """
 
-from contextlib import contextmanager
+import collections
 import importlib
 import pathlib
-import tempfile
 
 import pytest
 
@@ -24,73 +23,75 @@
 pq = pytest.importorskip("pyarrow.parquet")
 
 
-@contextmanager
-def create_catalog(catalog_name_in_pyiceberg_config=None):
-    # the catalog stores the full path of data files, so the catalog needs to be
-    # created dynamically, and not saved in pandas/tests/io/data as other formats
-    with tempfile.TemporaryDirectory("-pandas-iceberg.tmp") as catalog_path:
-        uri = f"sqlite:///{catalog_path}/catalog.sqlite"
-        warehouse = f"file://{catalog_path}"
-        catalog = pyiceberg_catalog.load_catalog(
-            catalog_name_in_pyiceberg_config or "default",
-            type="sql",
-            uri=uri,
-            warehouse=warehouse,
-        )
-        catalog.create_namespace("ns")
+Catalog = collections.namedtuple("name", "uri")
 
-        df = pq.read_table(
-            pathlib.Path(__file__).parent / "data" / "parquet" / "simple.parquet"
-        )
-        table = catalog.create_table("ns.my_table", schema=df.schema)
-        table.append(df)
 
-        if catalog_name_in_pyiceberg_config is not None:
-            config_path = pathlib.Path.home() / ".pyiceberg.yaml"
-            with open(config_path, "w", encoding="utf-8") as f:
-                f.write(f"""\
+@pytest.fixture
+def catalog(request, tmp_path, params=(None, "default", "pandas_tests")):
+    # the catalog stores the full path of data files, so the catalog needs to be
+    # created dynamically, and not saved in pandas/tests/io/data as other formats
+    catalog_path = tmp_path / "pandas-iceberg-catalog"
+    catalog_path.mkdir()
+    catalog_name = request.param
+    uri = f"sqlite:///{catalog_path}/catalog.sqlite"
+    warehouse = f"file://{catalog_path}"
+    catalog = pyiceberg_catalog.load_catalog(
+        catalog_name,
+        type="sql",
+        uri=uri,
+        warehouse=warehouse,
+    )
+    catalog.create_namespace("ns")
+
+    df = pq.read_table(
+        pathlib.Path(__file__).parent / "data" / "parquet" / "simple.parquet"
+    )
+    table = catalog.create_table("ns.my_table", schema=df.schema)
+    table.append(df)
+
+    if catalog_name is not None:
+        config_path = pathlib.Path.home() / ".pyiceberg.yaml"
+        with open(config_path, "w", encoding="utf-8") as f:
+            f.write(f"""\
 catalog:
-  {catalog_name_in_pyiceberg_config}:
+  {catalog_name}:
     type: sql
     uri: {uri}
     warehouse: {warehouse}""")
-        importlib.reload(pyiceberg_catalog)  # needed to reload the config file
 
-        try:
-            yield uri
-        finally:
-            if catalog_name_in_pyiceberg_config is not None:
-                config_path.unlink()
+    importlib.reload(pyiceberg_catalog)  # needed to reload the config file
+
+    yield Catalog(name=catalog_name, uri=uri)
+
+    if catalog_name is not None:
+        config_path.unlink()
 
 
 class TestIceberg:
-    def test_read(self):
+    def test_read(self, catalog):
         expected = pd.DataFrame(
             {
                 "A": [1, 2, 3],
                 "B": ["foo", "foo", "foo"],
             }
         )
-        with create_catalog() as catalog_uri:
-            result = read_iceberg(
-                "ns.my_table",
-                catalog_properties={"uri": catalog_uri},
-            )
+        result = read_iceberg(
+            "ns.my_table",
+            catalog_properties={"uri": catalog.uri},
+        )
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.parametrize("catalog_name", ["default", "pandas_tests"])
-    def test_read_by_catalog_name(self, catalog_name):
+    def test_read_by_catalog_name(self, catalog):
         expected = pd.DataFrame(
             {
                 "A": [1, 2, 3],
                 "B": ["foo", "foo", "foo"],
             }
         )
-        with create_catalog(catalog_name_in_pyiceberg_config=catalog_name):
-            result = read_iceberg(
-                "ns.my_table",
-                catalog_name=catalog_name,
-            )
+        result = read_iceberg(
+            "ns.my_table",
+            catalog_name=catalog.name,
+        )
         tm.assert_frame_equal(result, expected)
 
     def test_read_with_row_filter(self):
@@ -100,37 +101,34 @@ def test_read_with_row_filter(self):
                 "B": ["foo", "foo"],
             }
         )
-        with create_catalog() as catalog_uri:
-            result = read_iceberg(
-                "ns.my_table",
-                catalog_properties={"uri": catalog_uri},
-                row_filter="A > 1",
-            )
+        result = read_iceberg(
+            "ns.my_table",
+            catalog_properties={"uri": catalog.uri},
+            row_filter="A > 1",
+        )
         tm.assert_frame_equal(result, expected)
 
-    def test_read_with_case_sensitive(self):
+    def test_read_with_case_sensitive(self, catalog):
         expected = pd.DataFrame(
             {
                 "A": [1, 2, 3],
             }
         )
-        with create_catalog() as catalog_uri:
-            result = read_iceberg(
+        result = read_iceberg(
+            "ns.my_table",
+            catalog_properties={"uri": catalog.uri},
+            selected_fields=["a"],
+            case_sensitive=False,
+        )
+        tm.assert_frame_equal(result, expected)
+
+        with pytest.raises(ValueError, match="^Could not find column"):
+            read_iceberg(
                 "ns.my_table",
-                catalog_properties={"uri": catalog_uri},
+                catalog_properties={"uri": catalog.uri},
                 selected_fields=["a"],
-                case_sensitive=False,
+                case_sensitive=True,
             )
-        tm.assert_frame_equal(result, expected)
-
-        with create_catalog() as catalog_uri:
-            with pytest.raises(ValueError, match="^Could not find column"):
-                read_iceberg(
-                    "ns.my_table",
-                    catalog_properties={"uri": catalog_uri},
-                    selected_fields=["a"],
-                    case_sensitive=True,
-                )
 
     def test_read_with_limit(self):
         expected = pd.DataFrame(
@@ -139,10 +137,9 @@ def test_read_with_limit(self):
                 "B": ["foo", "foo"],
             }
         )
-        with create_catalog() as catalog_uri:
-            result = read_iceberg(
-                "ns.my_table",
-                catalog_properties={"uri": catalog_uri},
-                limit=2,
-            )
+        result = read_iceberg(
+            "ns.my_table",
+            catalog_properties={"uri": catalog.uri},
+            limit=2,
+        )
         tm.assert_frame_equal(result, expected)