diff --git a/packages/google-cloud-bigtable/conftest.py b/packages/google-cloud-bigtable/conftest.py new file mode 100644 index 000000000000..a95ecc3edb32 --- /dev/null +++ b/packages/google-cloud-bigtable/conftest.py @@ -0,0 +1,36 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Root pytest configuration for the package. + +Keeps ``pytest.ini`` portable across ``pytest-asyncio`` versions by registering +``asyncio_default_fixture_loop_scope`` and ``asyncio_default_test_loop_scope`` +when running under older ``pytest-asyncio`` (<0.23). +""" + + +def pytest_addoption(parser): + for name, help_text in ( + ( + "asyncio_default_fixture_loop_scope", + "default event-loop scope for asyncio fixtures (pytest-asyncio)", + ), + ( + "asyncio_default_test_loop_scope", + "default event-loop scope for asyncio tests (pytest-asyncio)", + ), + ): + try: + parser.addini(name, help_text, default=None) + except ValueError: + pass diff --git a/packages/google-cloud-bigtable/pytest.ini b/packages/google-cloud-bigtable/pytest.ini new file mode 100644 index 000000000000..1ef01ffdea10 --- /dev/null +++ b/packages/google-cloud-bigtable/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +asyncio_default_fixture_loop_scope = session +asyncio_default_test_loop_scope = session diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/README.md b/packages/google-cloud-bigtable/tests/system/data/accelerator/README.md new file mode 100644 index 000000000000..51c2c55a1e4b --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/README.md @@ -0,0 +1,31 @@ +# Cloud Bigtable Accelerator System Tests + +Shared environment setup, fixtures, and state/process helpers for live Cloud Bigtable Python accelerator system tests. Tests using this harness run against a real Cloud Bigtable instance and a real Go accelerator binary; emulator and binary-less environments skip cleanly. + +## Harness Components + +- [`conftest.py`](conftest.py): Registers the `accelerator` and `slow` pytest markers, provides `tmp_bin_dir`, and wraps `get_table()` on both sync and async data clients to inject `BIGTABLE_TEST_APP_PROFILE` (or `CBT_APP_PROFILE_ID`) when no `app_profile_id` is passed. +- [`_base_async.py`](_base_async.py) / [`_base_autogen.py`](_base_autogen.py): Define `AcceleratorTestBaseAsync` and its CrossSync sibling `AcceleratorTestBase`. They enforce live-instance and binary prerequisites per session, expose `client`, `accel_client`, `native_client`, `accel_table`, `native_table`, `default_table`, and `janitor` cleanup fixtures, and assert `table._accelerator_client` state via `assert_accelerator_active()` and `assert_native()`. +- [`_harness.py`](_harness.py): Provides the in-memory reference model (`ExpectedState`, `normalize_row`, `assert_rows_equivalent`), seeded workload generator (`RandomOps` with atomic multi-family/multi-column batches and `DELETE_AFTER_WRITE_GUARD_OPS = 512`), process leak/kill helpers (`ProcessIntrospector`, `force_kill_daemon`, `daemon_pid`, `pid_alive`), and binary/environment gates (`resolve_binary`, `require_binary_or_skip`, `require_real_bigtable_or_skip`). +- Package root [`pytest.ini`](../../../../pytest.ini) & [`conftest.py`](../../../../conftest.py): Set session-scoped `pytest-asyncio` loops and register the loop-scope ini options so older `pytest-asyncio` (<0.23) versions do not warn or fail. + +## Environment Variables + +| Variable | Required | Purpose | +| --- | --- | --- | +| `GOOGLE_CLOUD_PROJECT` (or `PROJECT_ID`) | Yes | Target GCP project; tests skip if unset. | +| `BIGTABLE_EMULATOR_HOST` | Must be unset | Emulator runs skip because the Go accelerator does not support the emulator. | +| `BIGTABLE_ACCELERATOR_BIN` | No | Path to a locally built Go accelerator binary; overrides the wheel-bundled binary lookup (`_daemon._resolve_binary_path()`). | +| `BIGTABLE_TEST_APP_PROFILE` (or `CBT_APP_PROFILE_ID`) | No | Default app profile ID injected into `get_table()` calls. | + +## Running Locally & Syncing + +Run pytest from `packages/google-cloud-bigtable/` once a compiled Go binary (or bundled wheel) and GCP project are available: + +```bash +BIGTABLE_ACCELERATOR_BIN=/path/to/accelerator \ +GOOGLE_CLOUD_PROJECT=my-project \ +pytest tests/system/data/accelerator -v +``` + +If you edit `_base_async.py`, regenerate `_base_autogen.py` from `packages/google-cloud-bigtable/` with `nox -s generate_sync`. diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py new file mode 100644 index 000000000000..58d482ea3866 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py new file mode 100644 index 000000000000..001b5d8e5612 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py @@ -0,0 +1,117 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import pytest + +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, RowMutationEntry + +from .. import SystemTestRunner +from . import _harness + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator._base_autogen" + + +@CrossSync.convert_class(sync_name="AcceleratorTestBase") +class AcceleratorTestBaseAsync(SystemTestRunner): + @pytest.fixture(scope="session", autouse=True) + def _require_accel_env(self): + _harness.require_real_bigtable_or_skip() + path = _harness.require_binary_or_skip() + orig_resolve = _daemon._resolve_binary_path + if os.environ.get(_harness.BIN_ENV_VAR): + _daemon._resolve_binary_path = lambda: path + try: + yield + finally: + _daemon._resolve_binary_path = orig_resolve + + def _project_id(self): + return os.getenv("GOOGLE_CLOUD_PROJECT") or os.getenv("PROJECT_ID") or None + + def _make_client(self, use_accelerator=None, **kwargs): + return CrossSync.DataClient( + project=self._project_id(), use_accelerator=use_accelerator, **kwargs + ) + + def assert_accelerator_active(self, table): + assert table._accelerator_client is not None, ( + "expected the accelerator to be active (use_accelerator=True) but the " + "target fell back to native. Check ADC principal / identity " + "verification and that the bundled daemon binary can start." + ) + + def assert_native(self, table): + assert table._accelerator_client is None, ( + "expected a native target (use_accelerator=False) but an accelerator " + "client was attached." + ) + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def client(self): + async with self._make_client() as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def accel_client(self): + async with self._make_client(use_accelerator=True) as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def native_client(self): + async with self._make_client(use_accelerator=False) as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def accel_table(self, accel_client, instance_id, table_id): + async with accel_client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def native_table(self, native_client, instance_id, table_id): + async with native_client.get_table(instance_id, table_id) as table: + self.assert_native(table) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def default_table(self, client, instance_id, table_id): + async with client.get_table(instance_id, table_id) as table: + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function") + async def janitor(self, native_table): + class _Janitor: + def __init__(self): + self.keys = set() + + def track(self, key): + self.keys.add(key) + return key + + j = _Janitor() + yield j + if j.keys: + entries = [RowMutationEntry(key, [DeleteAllFromRow()]) for key in j.keys] + await native_table.bulk_mutate_rows(entries) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py new file mode 100644 index 000000000000..399e15be0bc2 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# This file is automatically generated by CrossSync. Do not edit manually. + +import os + +import pytest + +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, RowMutationEntry + +from .. import SystemTestRunner +from . import _harness + + +class AcceleratorTestBase(SystemTestRunner): + @pytest.fixture(scope="session", autouse=True) + def _require_accel_env(self): + _harness.require_real_bigtable_or_skip() + path = _harness.require_binary_or_skip() + orig_resolve = _daemon._resolve_binary_path + if os.environ.get(_harness.BIN_ENV_VAR): + _daemon._resolve_binary_path = lambda: path + try: + yield + finally: + _daemon._resolve_binary_path = orig_resolve + + def _project_id(self): + return os.getenv("GOOGLE_CLOUD_PROJECT") or os.getenv("PROJECT_ID") or None + + def _make_client(self, use_accelerator=None, **kwargs): + return CrossSync._Sync_Impl.DataClient( + project=self._project_id(), use_accelerator=use_accelerator, **kwargs + ) + + def assert_accelerator_active(self, table): + assert table._accelerator_client is not None, ( + "expected the accelerator to be active (use_accelerator=True) but the target fell back to native. Check ADC principal / identity verification and that the bundled daemon binary can start." + ) + + def assert_native(self, table): + assert table._accelerator_client is None, ( + "expected a native target (use_accelerator=False) but an accelerator client was attached." + ) + + @pytest.fixture(scope="session") + def client(self): + with self._make_client() as client: + yield client + + @pytest.fixture(scope="session") + def accel_client(self): + with self._make_client(use_accelerator=True) as client: + yield client + + @pytest.fixture(scope="session") + def native_client(self): + with self._make_client(use_accelerator=False) as client: + yield client + + @pytest.fixture(scope="session") + def accel_table(self, accel_client, instance_id, table_id): + with accel_client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + yield table + + @pytest.fixture(scope="session") + def native_table(self, native_client, instance_id, table_id): + with native_client.get_table(instance_id, table_id) as table: + self.assert_native(table) + yield table + + @pytest.fixture(scope="session") + def default_table(self, client, instance_id, table_id): + with client.get_table(instance_id, table_id) as table: + yield table + + @pytest.fixture(scope="function") + def janitor(self, native_table): + class _Janitor: + def __init__(self): + self.keys = set() + + def track(self, key): + self.keys.add(key) + return key + + j = _Janitor() + yield j + if j.keys: + entries = [RowMutationEntry(key, [DeleteAllFromRow()]) for key in j.keys] + native_table.bulk_mutate_rows(entries) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py new file mode 100644 index 000000000000..d7058120014d --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py @@ -0,0 +1,412 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import gc +import glob +import os +import tempfile +import time +from dataclasses import dataclass +from typing import Any, Iterable, Sequence + +from google.cloud.bigtable.data import ( + AddToCell, + DeleteAllFromFamily, + DeleteAllFromRow, + DeleteRangeFromColumn, + SetCell, +) +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data.row import Row + +# Millisecond timestamp granularity (Bigtable default). +MS = 1000 + +# Column families created by the shared system-test conftest. +TEST_FAMILY = "test-family" +TEST_FAMILY_2 = "test-family-2" +TEST_AGGREGATE_FAMILY = "test-aggregate-family" +_AGGREGATE_VALUE_BYTES = 8 + +# Guard window (in ops) to avoid opposite-kind mutations on the same key within +# a retry straggler window. +DELETE_AFTER_WRITE_GUARD_OPS = 512 + +# Env var used by tests / runners to override the binary location. +BIN_ENV_VAR = "BIGTABLE_ACCELERATOR_BIN" + +Expected = tuple # (family: str, qualifier: bytes, timestamp_micros: int, value: bytes) + + +class ExpectedState: + def __init__(self) -> None: + self._rows: dict[bytes, dict[tuple[str, bytes], dict[int, Any]]] = {} + + def set_cell( + self, row_key: bytes, family: str, qualifier: bytes, value: bytes, ts: int + ) -> SetCell: + if ts % MS != 0: + raise ValueError( + "timestamps must be millisecond-granular (multiple of 1000)" + ) + col = self._rows.setdefault(row_key, {}).setdefault((family, qualifier), {}) + col[ts] = value + return SetCell(family, qualifier, value, timestamp_micros=ts) + + def add_to_cell( + self, row_key: bytes, qualifier: bytes, delta: int, ts: int + ) -> AddToCell: + if ts % MS != 0: + raise ValueError( + "timestamps must be millisecond-granular (multiple of 1000)" + ) + col = self._rows.setdefault(row_key, {}).setdefault( + (TEST_AGGREGATE_FAMILY, qualifier), {} + ) + col[ts] = int(col.get(ts, 0)) + delta + return AddToCell(TEST_AGGREGATE_FAMILY, qualifier, delta, timestamp_micros=ts) + + def delete_range_from_column( + self, + row_key: bytes, + family: str, + qualifier: bytes, + start: int | None = None, + end: int | None = None, + ) -> DeleteRangeFromColumn: + col = self._rows.get(row_key, {}).get((family, qualifier)) + if col is not None: + lo = 0 if start is None else start + for ts in list(col): + if ts >= lo and (end is None or ts < end): + del col[ts] + if not col: + self._rows[row_key].pop((family, qualifier), None) + if not self._rows[row_key]: + self._rows.pop(row_key, None) + return DeleteRangeFromColumn(family, qualifier, start, end) + + def delete_from_family(self, row_key: bytes, family: str) -> DeleteAllFromFamily: + row = self._rows.get(row_key) + if row is not None: + for key in [k for k in row if k[0] == family]: + del row[key] + if not row: + self._rows.pop(row_key, None) + return DeleteAllFromFamily(family) + + def delete_from_row(self, row_key: bytes) -> DeleteAllFromRow: + self._rows.pop(row_key, None) + return DeleteAllFromRow() + + def expected_cells(self, row_key: bytes) -> list[Expected]: + row = self._rows.get(row_key, {}) + out: list[Expected] = [] + for family, qualifier in sorted(row, key=lambda k: (k[0], k[1])): + for ts in sorted(row[(family, qualifier)], reverse=True): + value = row[(family, qualifier)][ts] + if family == TEST_AGGREGATE_FAMILY: + value = int(value).to_bytes( + _AGGREGATE_VALUE_BYTES, "big", signed=True + ) + out.append((family, qualifier, ts, value)) + return out + + def row_is_empty(self, row_key: bytes) -> bool: + return not self._rows.get(row_key) + + def keys(self) -> Iterable[bytes]: + return list(self._rows) + + +def normalize_row(row: Row | None) -> list[Expected]: + if row is None: + return [] + return [ + (cell.family, cell.qualifier, cell.timestamp_micros, cell.value) + for cell in row.cells + ] + + +def _fmt(cells: Sequence[Expected]) -> str: + return ( + "\n".join(f" {fam}:{qual!r}@{ts} = {val!r}" for (fam, qual, ts, val) in cells) + or " " + ) + + +def assert_rows_equivalent( + label_a: str, + cells_a: Sequence[Expected], + label_b: str, + cells_b: Sequence[Expected], +) -> None: + if list(cells_a) != list(cells_b): + raise AssertionError( + f"row mismatch between {label_a} and {label_b}:\n" + f"{label_a}:\n{_fmt(cells_a)}\n{label_b}:\n{_fmt(cells_b)}" + ) + + +def resolve_binary() -> str | None: + override = os.environ.get(BIN_ENV_VAR) + if override: + return override if os.path.isfile(override) else None + try: + return _daemon._resolve_binary_path() + except FileNotFoundError: + return None + + +def require_binary_or_skip() -> str: + import pytest + + path = resolve_binary() + if path is None: + pytest.skip( + "No accelerator daemon binary available " + f"(set {BIN_ENV_VAR} or install a wheel that bundles it)." + ) + return path + + +def require_real_bigtable_or_skip() -> None: + import pytest + from google.cloud.environment_vars import BIGTABLE_EMULATOR + + if os.environ.get(BIGTABLE_EMULATOR): + pytest.skip("accelerator is not supported against the emulator") + if not (os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("PROJECT_ID")): + pytest.skip("no GOOGLE_CLOUD_PROJECT set for live accelerator tests") + + +def accel_tempdirs() -> set[str]: + return set(glob.glob(os.path.join(tempfile.gettempdir(), "bt-accel-*"))) + + +@dataclass +class LeakSnapshot: + child_pids: set[int] + num_fds: int + tempdirs: set[str] + + +class ProcessIntrospector: + def __init__(self) -> None: + import psutil + + self._psutil = psutil + self._proc = psutil.Process(os.getpid()) + + def snapshot(self) -> LeakSnapshot: + children = set() + for child in self._proc.children(recursive=True): + try: + children.add(child.pid) + except self._psutil.Error: + pass + try: + num_fds = self._proc.num_fds() + except (self._psutil.Error, AttributeError): + num_fds = -1 + return LeakSnapshot( + child_pids=children, num_fds=num_fds, tempdirs=accel_tempdirs() + ) + + def assert_no_leaks( + self, + before: LeakSnapshot, + *, + fd_slack: int = 8, + label: str = "", + settle_timeout: float = 5.0, + ) -> None: + after = self.snapshot() + prefix = f"[{label}] " if label else "" + leaked_children = after.child_pids - before.child_pids + leaked_dirs = after.tempdirs - before.tempdirs + assert not leaked_children, ( + f"{prefix}leaked {len(leaked_children)} daemon subprocess(es): " + f"{sorted(leaked_children)}" + ) + assert not leaked_dirs, f"{prefix}leaked accel tempdirs: {sorted(leaked_dirs)}" + if before.num_fds < 0 or after.num_fds < 0: + return + gc.collect() + min_fds = after.num_fds + deadline = time.monotonic() + settle_timeout + while min_fds > before.num_fds + fd_slack and time.monotonic() < deadline: + time.sleep(0.1) + cur = self.snapshot().num_fds + if cur >= 0: + min_fds = min(min_fds, cur) + assert min_fds <= before.num_fds + fd_slack, ( + f"{prefix}fd count grew from {before.num_fds} to {min_fds} " + f"(slack {fd_slack}) and did not drain within {settle_timeout:g}s" + ) + + +def force_kill_daemon(daemon: Any) -> None: + proc = getattr(daemon, "_proc", None) + if proc is None: + return + try: + proc.kill() + except Exception: + pass + try: + proc.wait(timeout=2.0) + except Exception: + pass + + +def daemon_pid(table: Any) -> int | None: + daemon = getattr(table, "_accelerator_daemon", None) + if daemon is None: + return None + try: + return daemon.pid + except RuntimeError: + return None + + +def pid_alive(pid: int) -> bool: + import psutil + + try: + proc = psutil.Process(pid) + return proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + +class RandomOps: + def __init__( + self, + seed: int, + key_prefix: bytes = b"accel-", + *, + include_aggregate: bool = False, + ): + import random + + self._rand = random.Random(seed) + self._keyspace = [key_prefix + f"{i:08d}".encode() for i in range(64)] + self._families = [TEST_FAMILY, TEST_FAMILY_2] + self._qualifiers = [f"q{i}".encode() for i in range(self._rand.randint(2, 16))] + self._agg_qualifiers = [f"agg{i}".encode() for i in range(4)] + self._include_aggregate = include_aggregate + self._op_index = 0 + self._last_write_op: dict[bytes, int] = {} + self._last_delete_op: dict[bytes, int] = {} + + def key(self) -> bytes: + return self._rand.choice(self._keyspace) + + def value(self, max_len: int = 64) -> bytes: + n = self._rand.randint(0, max_len) + return bytes(self._rand.getrandbits(8) for _ in range(n)) + + def ms_timestamp(self) -> int: + return self._rand.randint(1, 2_000_000) * MS + + def family(self) -> str: + return self._rand.choice(self._families) + + def qualifier(self) -> bytes: + return self._rand.choice(self._qualifiers) + + def agg_qualifier(self) -> bytes: + return self._rand.choice(self._agg_qualifiers) + + def add_delta(self) -> int: + return self._rand.randint(-(2**20), 2**20) + + def _delete_target(self) -> tuple[str, bytes]: + if self._include_aggregate and self._rand.random() < 0.30: + return TEST_AGGREGATE_FAMILY, self.agg_qualifier() + return self.family(), self.qualifier() + + def _delete_family(self) -> str: + families = self._families + ( + [TEST_AGGREGATE_FAMILY] if self._include_aggregate else [] + ) + return self._rand.choice(families) + + def _build_multi_cf_mutations( + self, key: bytes, expected_state: ExpectedState + ) -> list[object]: + """Build an atomic write list touching multiple CFs and qualifiers.""" + mutations: list[object] = [] + for idx, fam in enumerate(self._families): + max_cols = min(3, len(self._qualifiers)) + min_cols = 2 if idx == 0 and max_cols >= 2 else 1 + num_cols = self._rand.randint(min_cols, max_cols) + for qual in self._rand.sample(self._qualifiers, num_cols): + mutations.append( + expected_state.set_cell( + key, fam, qual, self.value(), self.ms_timestamp() + ) + ) + if self._include_aggregate and self._rand.random() < 0.5: + mutations.append( + expected_state.add_to_cell( + key, self.agg_qualifier(), self.add_delta(), self.ms_timestamp() + ) + ) + self._rand.shuffle(mutations) + return mutations + + def build_mutation(self, expected_state: ExpectedState) -> tuple[bytes, object]: + self._op_index += 1 + roll = self._rand.random() + key = self.key() + last_write = self._last_write_op.get(key) + last_delete = self._last_delete_op.get(key) + recently_written = ( + last_write is not None + and self._op_index - last_write <= DELETE_AFTER_WRITE_GUARD_OPS + ) + recently_deleted = ( + last_delete is not None + and self._op_index - last_delete <= DELETE_AFTER_WRITE_GUARD_OPS + ) + if recently_deleted and not recently_written: + self._last_delete_op[key] = self._op_index + return key, expected_state.delete_from_row(key) + if self._include_aggregate and roll < 0.18: + self._last_write_op[key] = self._op_index + return key, expected_state.add_to_cell( + key, self.agg_qualifier(), self.add_delta(), self.ms_timestamp() + ) + if roll < 0.70 or recently_written: + self._last_write_op[key] = self._op_index + if self._rand.random() < 0.35: + return key, self._build_multi_cf_mutations(key, expected_state) + return key, expected_state.set_cell( + key, self.family(), self.qualifier(), self.value(), self.ms_timestamp() + ) + self._last_delete_op[key] = self._op_index + if roll < 0.82: + a = self._rand.choice([None, self.ms_timestamp()]) + b = self._rand.choice([None, self.ms_timestamp()]) + if a is not None and b is not None and a > b: + a, b = b, a + fam, qual = self._delete_target() + return key, expected_state.delete_range_from_column(key, fam, qual, a, b) + if roll < 0.92: + return key, expected_state.delete_from_family(key, self._delete_family()) + return key, expected_state.delete_from_row(key) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py new file mode 100644 index 000000000000..f7b5cfe1549a --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py @@ -0,0 +1,68 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +import os +import tempfile + +import pytest + + +@pytest.fixture(scope="session", autouse=True) +def _default_app_profile(): + """Default ``get_table`` to ``BIGTABLE_TEST_APP_PROFILE`` when set.""" + profile = os.getenv("BIGTABLE_TEST_APP_PROFILE") or os.getenv("CBT_APP_PROFILE_ID") + if not profile: + yield + return + + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + from google.cloud.bigtable.data._sync_autogen.client import BigtableDataClient + + patched = [] + for cls in (BigtableDataClientAsync, BigtableDataClient): + original = cls.get_table + + def make_wrapper(orig): + @functools.wraps(orig) + def get_table(self, instance_id, table_id, *args, **kwargs): + if not args and "app_profile_id" not in kwargs: + kwargs["app_profile_id"] = profile + return orig(self, instance_id, table_id, *args, **kwargs) + + return get_table + + cls.get_table = make_wrapper(original) + patched.append((cls, original)) + + yield + + for cls, original in patched: + cls.get_table = original + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "accelerator: integration tests for the Cloud Bigtable accelerator", + ) + config.addinivalue_line( + "markers", "slow: longer-running accelerator tests (leak loops, stress)" + ) + + +@pytest.fixture +def tmp_bin_dir(): + with tempfile.TemporaryDirectory(prefix="accel-test-bin-") as d: + yield d