diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 7f6e8e507..78cb7a991 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -342,7 +342,7 @@ def _revise_request_rowset( _RowSetComplete: if there are no rows left to process after the revision """ # if user is doing a whole table scan, start a new one with the last seen key - if row_set is None or (not row_set.row_ranges and row_set.row_keys is not None): + if row_set is None or (not row_set.row_ranges and not row_set.row_keys): last_seen = last_seen_row_key return RowSetPB(row_ranges=[RowRangePB(start_key_open=last_seen)]) # remove seen keys from user-specific key list diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index 4e7797c6d..2bf8688fd 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -98,19 +98,31 @@ def test_ctor(self): (["d", "c", "b", "a"], "b", ["d", "c"]), ], ) - def test_revise_request_rowset_keys(self, in_keys, last_key, expected): + @pytest.mark.parametrize("with_range", [True, False]) + def test_revise_request_rowset_keys_with_range( + self, in_keys, last_key, expected, with_range + ): from google.cloud.bigtable_v2.types import RowSet as RowSetPB from google.cloud.bigtable_v2.types import RowRange as RowRangePB + from google.cloud.bigtable.data.exceptions import _RowSetComplete in_keys = [key.encode("utf-8") for key in in_keys] expected = [key.encode("utf-8") for key in expected] last_key = last_key.encode("utf-8") - sample_range = RowRangePB(start_key_open=last_key) - row_set = RowSetPB(row_keys=in_keys, row_ranges=[sample_range]) - revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised.row_keys == expected - assert revised.row_ranges == [sample_range] + if with_range: + sample_range = [RowRangePB(start_key_open=last_key)] + else: + sample_range = [] + row_set = RowSetPB(row_keys=in_keys, row_ranges=sample_range) + if not with_range and expected == []: + # expect exception if we are revising to an empty rowset + with pytest.raises(_RowSetComplete): + self._get_target_class()._revise_request_rowset(row_set, last_key) + else: + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised.row_keys == expected + assert revised.row_ranges == sample_range @pytest.mark.parametrize( "in_ranges,last_key,expected", @@ -157,9 +169,13 @@ def test_revise_request_rowset_keys(self, in_keys, last_key, expected): ), ], ) - def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): + @pytest.mark.parametrize("with_key", [True, False]) + def test_revise_request_rowset_ranges( + self, in_ranges, last_key, expected, with_key + ): from google.cloud.bigtable_v2.types import RowSet as RowSetPB from google.cloud.bigtable_v2.types import RowRange as RowRangePB + from google.cloud.bigtable.data.exceptions import _RowSetComplete # convert to protobuf next_key = (last_key + "a").encode("utf-8") @@ -172,10 +188,20 @@ def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) for r in expected ] - row_set = RowSetPB(row_ranges=in_ranges, row_keys=[next_key]) - revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised.row_keys == [next_key] - assert revised.row_ranges == expected + if with_key: + row_keys = [next_key] + else: + row_keys = [] + + row_set = RowSetPB(row_ranges=in_ranges, row_keys=row_keys) + if not with_key and expected == []: + # expect exception if we are revising to an empty rowset + with pytest.raises(_RowSetComplete): + self._get_target_class()._revise_request_rowset(row_set, last_key) + else: + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised.row_keys == row_keys + assert revised.row_ranges == expected @pytest.mark.parametrize("last_key", ["a", "b", "c"]) def test_revise_request_full_table(self, last_key):