test(pyamber): cover IcebergDocument uri, count and clear paths

amber/src/test/python/core/storage/iceberg/test_iceberg_document.py

@@ -22,10 +22,21 @@
 import uuid
 from concurrent.futures import as_completed
 from concurrent.futures.thread import ThreadPoolExecutor
+from unittest.mock import MagicMock, Mock, patch
+from urllib.parse import urlparse
+
+from pyiceberg import types as iceberg_types
+from pyiceberg.schema import Schema as IcebergSchema
 
 from core.models import Schema, Tuple
 from core.models.state import State
 from core.storage.document_factory import DocumentFactory
+from core.storage.iceberg import iceberg_document
+from core.storage.iceberg.iceberg_document import IcebergDocument
+from core.storage.iceberg.iceberg_utils import (
+    amber_tuples_to_arrow_table,
+    arrow_table_to_amber_tuples,
+)
 from core.storage.storage_config import StorageConfig
 from core.storage.vfs_uri_factory import VFSURIFactory
 from proto.org.apache.texera.amber.core import (
@@ -411,3 +422,183 @@ def test_multiple_states_materialize_as_rows_in_one_table(self):
             key=lambda state: state["loop_counter"],
         )
         assert actual_states == states
+
+
+class TestIcebergDocumentWithMockCatalog:
+    """
+    The catalog-facing paths of IcebergDocument that need no catalog service:
+    the table-location lookup, the "table is absent" arms of clear() and
+    get_count(), the locks clear() and the read path take, what each read entry
+    point hands the iterator, and both sides of the iterator's seek guard. The
+    catalog is a mock, so unlike TestIcebergDocument above these run on any host
+    without a live Iceberg catalog, and they mutate no cached catalog state.
+    """
+
+    @pytest.fixture
+    def iceberg_schema(self):
+        return IcebergSchema(
+            iceberg_types.NestedField(
+                field_id=1,
+                name="col-int",
+                field_type=iceberg_types.IntegerType(),
+                required=False,
+            )
+        )
+
+    @pytest.fixture
+    def document(self, iceberg_schema):
+        """
+        An IcebergDocument for `ns.tbl` whose catalog is a mock. `get_instance`
+        is patched only for the duration of the construction, so the real
+        per-warehouse catalog cache is never touched.
+        """
+        catalog = Mock()
+        with patch.object(
+            iceberg_document.IcebergCatalogInstance,
+            "get_instance",
+            return_value=catalog,
+        ):
+            return IcebergDocument(
+                "ns",
+                "tbl",
+                iceberg_schema,
+                amber_tuples_to_arrow_table,
+                arrow_table_to_amber_tuples,
+            )
+
+    def test_get_uri_returns_the_parsed_table_location(self, document):
+        # The location is unique per run and the expectation is derived from it,
+        # so a get_uri that returned a fixed URI instead of parsing the loaded
+        # table's own location could not accidentally match.
+        location = f"file:///warehouse/{uuid.uuid4().hex}/ns.db/tbl"
+        table = Mock()
+        table.location.return_value = location
+
+        with patch.object(
+            iceberg_document, "load_table_metadata", return_value=table
+        ) as load_table_metadata:
+            uri = document.get_uri()
+
+        expected = urlparse(location)
+        assert (uri.scheme, uri.path) == (expected.scheme, expected.path)
+        table.location.assert_called_once_with()
+        assert load_table_metadata.call_args.args == (document.catalog, "ns", "tbl")
+
+    def test_get_uri_rejects_a_table_that_does_not_exist(self, document):
+        with patch.object(iceberg_document, "load_table_metadata", return_value=None):
+            with pytest.raises(Exception, match=r"table ns\.tbl doesn't exist\."):
+                document.get_uri()
+
+    def test_clear_drops_a_table_that_exists(self, document):
+        document.catalog.table_exists.return_value = True
+
+        document.clear()
+
+        document.catalog.drop_table.assert_called_once_with("ns.tbl")
+
+    def test_clear_leaves_an_absent_table_alone(self, document):
+        document.catalog.table_exists.return_value = False
+
+        document.clear()
+
+        document.catalog.table_exists.assert_called_once_with("ns.tbl")
+        document.catalog.drop_table.assert_not_called()
+
+    def test_get_count_is_zero_when_the_table_does_not_exist(self, document):
+        with patch.object(
+            iceberg_document, "load_table_metadata", return_value=None
+        ) as load_table_metadata:
+            assert document.get_count() == 0
+
+        assert load_table_metadata.call_args.args == (document.catalog, "ns", "tbl")
+
+    def test_a_negative_offset_is_rejected_rather_than_read_as_zero(self, document):
+        """
+        IcebergIterator guards its file seek against having already skipped past
+        `from_index`. The only way the guard can fire is a negative offset, since
+        the skip counter is still 0 when the seek generator first runs.
+
+        Neither the guard's message ("seek operation should not be called", which
+        describes a re-entrant seek rather than a bad argument) nor its exception
+        type is pinned: rejecting a negative offset as a ValueError would be the
+        better behaviour, so accepting either type here keeps that fix open while
+        still requiring that the offset is rejected rather than read as zero.
+        """
+        iterator = document.get_after(-1)
+
+        with pytest.raises((RuntimeError, ValueError)):
+            next(iterator)
+
+    @pytest.mark.parametrize("offset", [0, 5])
+    def test_a_non_negative_offset_does_not_trip_the_seek_guard(self, document, offset):
+        """
+        The other side of the seek guard's boundary: a legal offset must reach
+        the table lookup and then end the iteration cleanly, not raise. Without
+        this, the guard's comparison is unconstrained on hosts that cannot run
+        TestIcebergDocument above.
+        """
+        with patch.object(
+            iceberg_document, "load_table_metadata", return_value=None
+        ) as load_table_metadata:
+            with pytest.raises(StopIteration):
+                next(document.get_after(offset))
+
+        assert load_table_metadata.call_args.args == (document.catalog, "ns", "tbl")
+
+    @pytest.mark.parametrize(
+        "read, from_index, until_index, total",
+        [
+            (lambda document: document.get(), 0, None, float("inf")),
+            (lambda document: document.get_range(3, 7), 3, 7, 4),
+            (lambda document: document.get_after(4), 4, None, float("inf")),
+        ],
+    )
+    def test_the_read_entry_points_delegate_to_the_iterator(
+        self, document, read, from_index, until_index, total
+    ):
+        """
+        Each read entry point hands a specific [from, until) range -- and the
+        document's own catalog, table identity, schema and deserde -- to the
+        iterator. The seek generator's body does not run at construction, so
+        this needs no catalog.
+
+        `total` is spelled out per case rather than recomputed from the range,
+        so the expectation does not restate the production formula.
+        """
+        iterator = read(document)
+
+        assert (iterator.from_index, iterator.until_index) == (from_index, until_index)
+        assert iterator.total_records_to_return == total
+        assert (iterator.table_namespace, iterator.table_name) == ("ns", "tbl")
+        assert iterator.catalog is document.catalog
+        assert iterator.table_schema is document.table_schema
+        assert iterator.deserde is document.deserde
+
+    def test_the_read_path_takes_the_shared_read_lock(self, document):
+        """
+        The counterpart of test_clear_takes_the_write_lock: reads must take the
+        shared read lock, so that concurrent reads are not serialised behind
+        each other. MagicMock (not Mock) is required: the lock is used as a
+        context manager.
+        """
+        document.lock = MagicMock()
+
+        document.get()
+
+        document.lock.gen_rlock.assert_called_once_with()
+        document.lock.gen_wlock.assert_not_called()
+
+    def test_clear_takes_the_write_lock(self, document):
+        """
+        clear() drops the table, so it must hold the write lock rather than the
+        shared read lock the readers take. MagicMock (not Mock) is required: the
+        lock is used as a context manager.
+        """
+        document.catalog.table_exists.return_value = True
+        document.lock = MagicMock()
+
+        document.clear()
+
+        document.lock.gen_wlock.assert_called_once_with()
+        document.lock.gen_rlock.assert_not_called()
+        document.catalog.drop_table.assert_called_once_with("ns.tbl")