fix(ci): run the real Python test suite (3 tests → 182)

.codex/skills/ci-pr-helper/scripts/run_ci_checks.sh

@@ -15,12 +15,12 @@ fi
 
 cd "$root"
 
-cargo fmt --manifest-path rust/lance-context/Cargo.toml -- --check
-cargo clippy --manifest-path rust/lance-context/Cargo.toml --all-targets -- -D warnings
-cargo test --manifest-path rust/lance-context/Cargo.toml
+cargo fmt --all -- --check
+cargo clippy --workspace --all-targets -- -D warnings
+cargo test --workspace --all-targets
 
 cd "$root/python"
 uv run pytest
-uv run ruff format --check python/
-uv run ruff check python/
+uv run ruff format --check .
+uv run ruff check .
 uv run pyright

.github/workflows/python-test.yml

@@ -140,7 +140,7 @@ jobs:
         run: |
           uv venv
           source .venv/bin/activate
-          uv pip install pytest
+          uv pip install pytest pytest-asyncio
       - name: Download wheel
         uses: actions/download-artifact@v4
         with:
@@ -151,20 +151,25 @@ jobs:
         run: |
           source .venv/bin/activate
           WHEEL=$(ls dist/*.whl)
-          uv pip install "${WHEEL}[lance-python]"
+          uv pip install "${WHEEL}[lance-python,tests]"
 
       - name: Run tests
         working-directory: python
         run: |
           source .venv/bin/activate
-          pytest python/tests/ -v
+          pytest -v
       - name: Run doctests
         working-directory: python
         run: |
           source .venv/bin/activate
-          if [ -f python/lance_context/__init__.py ]; then
-            python -m doctest python/lance_context/__init__.py || echo "No doctests found"
-          fi
+          # Import as a module, not a file path: `python -m doctest <path>`
+          # imports by filename, which breaks the package-relative imports in
+          # __init__.py. Runs against the installed wheel and fails on any
+          # failing doctest.
+          python -c "import doctest, sys, lance_context; \
+            r = doctest.testmod(lance_context); \
+            print(r); \
+            sys.exit(1 if r.failed else 0)"
 
   lint:
     runs-on: ubuntu-24.04
@@ -198,12 +203,12 @@ jobs:
         working-directory: python
         run: |
           source .venv/bin/activate
-          ruff format --check python/
+          ruff format --check .
       - name: Run ruff lint
         working-directory: python
         run: |
           source .venv/bin/activate
-          ruff check python/
+          ruff check .
       - name: Run pyright type check
         working-directory: python
         run: |

crates/lance-context/src/unified_rollout.rs

@@ -64,6 +64,27 @@ impl RolloutStore {
             .map_err(|e| ContextError::Internal(e.to_string()))?;
         Ok(Self::Remote(store))
     }
+
+    /// Seal the local MemWAL memtable so rows written by [`RolloutStoreApi::add`]
+    /// become visible to subsequent reads on this handle.
+    ///
+    /// `add` is durable on return but *not* visible on return: visibility is
+    /// driven by a periodic sweeper in the server. An embedded caller has no
+    /// sweeper, so without an explicit `flush` a write-then-read sequence
+    /// returns nothing. Call this before reading back rows you just wrote.
+    ///
+    /// A no-op for `Remote` stores, where the server owns flush scheduling and
+    /// per-request `?flush=true` provides read-your-write.
+    pub async fn flush(&self) -> Result<(), ContextError> {
+        match self {
+            RolloutStore::Local(s) => s
+                .flush()
+                .await
+                .map_err(|e| ContextError::Internal(e.to_string())),
+            #[cfg(feature = "remote")]
+            RolloutStore::Remote(_) => Ok(()),
+        }
+    }
 }
 
 macro_rules! dispatch_mut {

python/pyproject.toml

@@ -60,6 +60,16 @@ tests = [
 ]
 dev = ["ruff", "pyright"]
 
+[tool.pytest.ini_options]
+# `testpaths` is relative to this file (python/), so this resolves to
+# python/tests/ -- the real suite. Without it, running pytest from python/
+# with a bare `pytest` argument could collect python/python/tests/ instead.
+testpaths = ["tests"]
+asyncio_mode = "strict"
+# Fail on unregistered marks so a typo'd @pytest.mark.asyncio cannot silently
+# turn an async test into a no-op pass.
+addopts = "--strict-markers --strict-config"
+
 [tool.ruff]
 lint.select = ["F", "E", "W", "I", "G", "TCH", "PERF", "B019"]
 

python/python/lance_context/api.py

@@ -2579,6 +2579,19 @@ def add_one(self, **fields: Any) -> dict[str, Any]:
         """Append a single record given as keyword arguments."""
         return self.add(fields)
 
+    def flush(self) -> None:
+        """Make previously added rows visible to subsequent reads.
+
+        ``add`` is durable on return but *not* visible on return: rows land in
+        a write-ahead log and only reach a readable generation when the
+        memtable is sealed. A deployed server does this on a timer; an embedded
+        store has no such timer, so a write-then-read sequence returns nothing
+        until you call this.
+
+        No-op for stores connected to a remote server.
+        """
+        self._sync.flush()
+
     def list(
         self,
         limit: int | None = None,
@@ -2672,6 +2685,17 @@ async def add_one(self, **fields: Any) -> dict[str, Any]:
         """Append a single record given as keyword arguments."""
         return await self.add(fields)
 
+    async def flush(self) -> None:
+        """Make previously added rows visible to subsequent reads.
+
+        For an embedded store this seals the local memtable. For a store
+        connected to a remote server this is a no-op: the server flushes on its
+        own interval (``ROLLOUT_FLUSH_INTERVAL_SECS``), which bounds how long a
+        just-written row stays invisible.
+        """
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self._sync.flush)
+
     async def list(
         self,
         limit: int | None = None,

python/src/lib.rs

@@ -2453,6 +2453,13 @@ impl RolloutStore {
         self.store.version()
     }
 
+    /// Seal the local MemWAL memtable so previously added rows become visible
+    /// to subsequent reads. No-op for remote stores.
+    fn flush(&self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| self.runtime.block_on(self.store.flush()))
+            .map_err(to_py_err)
+    }
+
     /// Append rollout rows given as a JSON array of `AddRolloutRequest` objects.
     /// Returns a dict `{version, ids, count}`.
     fn add(&mut self, py: Python<'_>, records_json: &str) -> PyResult<PyObject> {

python/tests/test_persistence.py

@@ -601,6 +601,23 @@ def test_time_travel_checkout(tmp_path: Path) -> None:
     ]
 
 
+_S3_MEMWAL_XFAIL = pytest.mark.xfail(
+    strict=True,
+    reason=(
+        "Upstream: lance 7.0.0 mem_wal_writer builds its object store with "
+        "ObjectStore::from_uri(base_uri), dropping the dataset's "
+        "storage_options (lance/src/dataset/mem_wal/api.rs:607). The WAL "
+        "writer therefore ignores aws_endpoint_url and targets real AWS, so "
+        "`add` fails with 'bucket not found' against any custom-endpoint S3 "
+        "(moto, MinIO, Ceph, R2). `create` works because it goes through the "
+        "options-aware path. Needs from_uri_and_params upstream; there is no "
+        "local workaround since ShardWriterConfig cannot carry the options. "
+        "Remove this marker once lance is bumped past the fix."
+    ),
+)
+
+
+@_S3_MEMWAL_XFAIL
 def test_s3_round_trip_with_storage_options(moto_endpoint: str, s3_client) -> None:
     """Canonical path: generic storage_options dict (aligns with lance/lance-graph)."""
     bucket = f"context-{uuid.uuid4().hex}"
@@ -619,6 +636,7 @@ def test_s3_round_trip_with_storage_options(moto_endpoint: str, s3_client) -> No
     assert ctx.entries() == 2
 
 
+@_S3_MEMWAL_XFAIL
 def test_s3_deprecated_aws_kwargs_still_work(moto_endpoint: str, s3_client) -> None:
     """AWS kwargs keep working (back-compat) and emit a DeprecationWarning."""
     bucket = f"context-{uuid.uuid4().hex}"

python/tests/test_rollout.py

@@ -38,6 +38,9 @@ def test_add_dict_list_get(store_uri):
     assert resp["count"] == 1
     assert resp["ids"] == ["row-0"]
 
+    # `add` is durable on return but not visible until the memtable is sealed.
+    store.flush()
+
     rows = store.list()
     assert len(rows) == 1
     assert rows[0]["id"] == "row-0"
@@ -60,6 +63,7 @@ def test_add_many_and_add_one(store_uri):
         ]
     )
     store.add_one(id="row-3", rollout_id="traj-1", reward=3.0)
+    store.flush()
 
     rows = store.list()
     assert {r["id"] for r in rows} == {"row-0", "row-1", "row-2", "row-3"}
@@ -75,6 +79,7 @@ def test_id_autogenerated_when_omitted(store_uri):
     (generated,) = resp["ids"]
     # Round-trips as a valid UUID and is fetchable by the generated id.
     assert uuid.UUID(generated).version == 4
+    store.flush()
     row = store.get(generated)
     assert row is not None and row["id"] == generated
 
@@ -97,6 +102,7 @@ def test_binary_payload_roundtrip(store_uri):
             "payload_size": len(blob),
         }
     )
+    store.flush()
 
     # list/get project the blob column out (cheap metadata scans).
     row = store.get("art-0")
@@ -121,6 +127,9 @@ def test_empty_add_rejected(store_uri):
 def test_reopen_sees_prior_rows(store_uri):
     store = RolloutStore.open(store_uri)
     store.add({"id": "row-0", "rollout_id": "traj-1"})
+    # Seal before dropping: an unsealed memtable is durable in the WAL but is
+    # not replayed into a readable generation by a fresh reader.
+    store.flush()
     del store
 
     reopened = RolloutStore.open(store_uri)
@@ -160,6 +169,7 @@ def test_list_filters_before_pagination(store_uri):
             },
         ]
     )
+    store.flush()
 
     rows = store.list(
         filters={"policy_version": "ckpt-2", "include_in_training": False}
@@ -190,6 +200,7 @@ def test_get_trajectory_orders_rows(store_uri):
             {"id": "row-b", "rollout_id": "target", "sequence_order": 0},
         ]
     )
+    store.flush()
 
     rows = store.get_trajectory("target")
     assert [row["id"] for row in rows] == ["row-b", "row-a"]

python/tests/test_rollout_remote.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import socket
 import subprocess
 import tempfile
@@ -43,12 +44,35 @@ def _wait_for_health(base_url: str, timeout: float = 30.0) -> None:
     raise RuntimeError(f"server did not become healthy at {url}")
 
 
+async def _eventually(fn, predicate, timeout: float = 15.0):
+    """Poll `fn` until `predicate` holds, or fail after `timeout`.
+
+    `add` is durable on return but not visible until the server's sweeper seals
+    the memtable, so a read immediately after a write legitimately returns
+    nothing. Polling asserts the row *arrives* without pinning the test to the
+    flush interval.
+    """
+    deadline = time.time() + timeout
+    last = None
+    while time.time() < deadline:
+        last = await fn()
+        if predicate(last):
+            return last
+        await asyncio.sleep(0.1)
+    raise AssertionError(f"condition not met within {timeout}s; last value: {last!r}")
+
+
 @pytest.fixture()
 def server():
     if not _SERVER_BIN.exists():
         pytest.skip(f"server binary not built at {_SERVER_BIN}")
     port = _free_port()
     with tempfile.TemporaryDirectory() as data_dir:
+        # Rows are durable on `add` but only become visible when the server's
+        # sweeper seals the memtable. The 30s production default would make
+        # every write-then-assert below hang; 1s keeps the tests honest about
+        # the async-visibility contract without waiting on it.
+        env = {**os.environ, "ROLLOUT_FLUSH_INTERVAL_SECS": "1"}
         proc = subprocess.Popen(
             [
                 str(_SERVER_BIN),
@@ -59,6 +83,7 @@ def server():
                 "--data-dir",
                 data_dir,
             ],
+            env=env,
             stdout=subprocess.DEVNULL,
             stderr=subprocess.DEVNULL,
         )
@@ -101,7 +126,7 @@ async def run():
         )
         assert resp["count"] == 2
 
-        rows = await store.list()
+        rows = await _eventually(store.list, lambda r: len(r) == 2)
         assert {r["id"] for r in rows} == {"row-0", "row-1"}
 
         one = await store.get("row-0")
@@ -125,7 +150,7 @@ async def run():
         # A second connection sees the first's flushed write (durable, no
         # read affinity).
         reader = await AsyncRolloutStore.connect(server, "rl-run-2")
-        rows = await reader.list()
+        rows = await _eventually(reader.list, lambda r: len(r) == 1)
         assert [r["id"] for r in rows] == ["only"]
 
     asyncio.run(run())
@@ -153,8 +178,11 @@ async def run():
             ]
         )
 
-        rows = await store.list(
-            filters={"policy_version": "ckpt-7", "include_in_training": True}
+        rows = await _eventually(
+            lambda: store.list(
+                filters={"policy_version": "ckpt-7", "include_in_training": True}
+            ),
+            lambda r: len(r) == 1,
         )
         assert [row["id"] for row in rows] == ["row-7"]
 
@@ -172,7 +200,9 @@ async def run():
             ]
         )
 
-        rows = await store.get_trajectory("target")
+        rows = await _eventually(
+            lambda: store.get_trajectory("target"), lambda r: len(r) == 2
+        )
         assert [row["id"] for row in rows] == ["row-b", "row-a"]
 
     asyncio.run(run())

python/tests/test_search.py

@@ -31,6 +31,8 @@ def __init__(self) -> None:
         ] = []
         self.list_calls: list[tuple[int | None, int | None, str | None]] = []
         self.list_lifecycle_calls: list[tuple[bool, bool]] = []
+        self.list_projection_calls: list[tuple[bool, bool]] = []
+        self.search_projection_calls: list[tuple[bool, bool]] = []
         self.related_calls: list[tuple[str, str | None, int | None, bool, bool]] = []
         self.get_calls: list[tuple[str | None, str | None]] = []
         self.delete_calls: list[tuple[str | None, str | None]] = []
@@ -274,10 +276,13 @@ def search(
         include_expired: bool = False,
         include_retired: bool = False,
         include_relationships: bool = False,
+        include_binary: bool = True,
+        include_embedding: bool = True,
     ):
         self.search_calls.append((vector, limit, filters_json))
         self.search_lifecycle_calls.append((include_expired, include_retired))
         self.search_relationship_calls.append(include_relationships)
+        self.search_projection_calls.append((include_binary, include_embedding))
         hit = {
             "id": "rec-1",
             "external_id": "source-1",
@@ -362,9 +367,12 @@ def list(
         filters_json: str | None,
         include_expired: bool = False,
         include_retired: bool = False,
+        include_binary: bool = True,
+        include_embedding: bool = True,
     ):
         self.list_calls.append((limit, offset, filters_json))
         self.list_lifecycle_calls.append((include_expired, include_retired))
+        self.list_projection_calls.append((include_binary, include_embedding))
         return [
             {
                 "id": "rec-1",