From 03e0dc5a21042bc9d634f81c716598abbd3bbe74 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 7 Oct 2026 14:24:41 +0000 Subject: [PATCH] [Python] Fix BigtableIO write test mocks for google-cloud-bigtable batcher delegation In google-cloud-bigtable >= 2.48.0 (PR #18200), MutationsBatcher delegates directly to the underlying data client batcher via table._table_impl.mutations_batcher instead of calling Table.mutate_rows. When using a mock Bigtable client in tests, table._table_impl was an unconfigured mock returning a default batcher whose close() was a no-op, causing Table.mutate_rows to be bypassed and leaving exceptions and batch completion callbacks un-triggered. In TestWriteBigTable.setUp, configure the mock batcher on table._table_impl.mutations_batcher to invoke Table.mutate_rows and forward the results to the registered batch completion callback. This ensures tests patching Table.mutate_rows (test_write_metrics, test_write_batch_error_surfaces_from_async_flush, and test_write_batch_error_surfaces_from_buffered_rows) work seamlessly across both legacy and modern client versions while keeping production code untouched. --- .../apache_beam/io/gcp/bigtableio_test.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 4aba9a4062fd..d62ae7f9ab0b 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -281,6 +281,31 @@ def setUp(self): instance = Instance(self._INSTANCE_ID, client) self.table = Table(self._TABLE_ID, instance) + # In google-cloud-bigtable >= 2.48.0, MutationsBatcher delegates to the + # underlying data client batcher instead of invoking Table.mutate_rows. + # Wire the mock data client batcher to surface Table.mutate_rows calls, + # exceptions, and callbacks so tests patching Table.mutate_rows work + # seamlessly. + def mock_data_client_batcher(*args, **kwargs): + batcher = MagicMock() + + def mock_close(): + res = self.table.mutate_rows([]) + cb = getattr(batcher, '_user_batch_completed_callback', None) + if cb and res and isinstance(res, (list, tuple)): + cb(res) + + batcher.close = MagicMock(side_effect=mock_close) + return batcher + + if hasattr(self.table, '_table_impl'): + self.table._table_impl.mutations_batcher = MagicMock( + side_effect=mock_data_client_batcher) + client._veneer_data_client.get_table.return_value.mutations_batcher = ( + MagicMock(side_effect=mock_data_client_batcher)) + client.data_client.mutations_batcher = MagicMock( + side_effect=mock_data_client_batcher) + def test_write(self): direct_rows = [self.generate_row(i) for i in range(5)] # TODO(https://github.com/apache/beam/issues/34549): This test relies on