use cancelled error to test non-retryable methods

Author: daniel-sanche

google/cloud/bigtable/data/_metrics/data_model.py

@@ -280,7 +280,7 @@ def _parse_response_metadata_blob(blob: bytes) -> Tuple[str, str] | None:
             # failed to parse metadata
             return None
 
-    def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
+    def end_attempt_with_status(self, status: StatusCode | BaseException) -> None:
         """
         Called to mark the end of an attempt for the operation.
 
@@ -297,7 +297,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
             return self._handle_error(
                 INVALID_STATE_ERROR.format("end_attempt_with_status", self.state)
             )
-        if isinstance(status, Exception):
+        if isinstance(status, BaseException):
             status = self._exc_to_status(status)
         complete_attempt = CompletedAttemptMetric(
             duration_ns=time.monotonic_ns() - self.active_attempt.start_time_ns,
@@ -312,7 +312,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
         for handler in self.handlers:
             handler.on_attempt_complete(complete_attempt, self)
 
-    def end_with_status(self, status: StatusCode | Exception) -> None:
+    def end_with_status(self, status: StatusCode | BaseException) -> None:
         """
         Called to mark the end of the operation. If there is an active attempt,
         end_attempt_with_status will be called with the same status.
@@ -329,7 +329,7 @@ def end_with_status(self, status: StatusCode | Exception) -> None:
                 INVALID_STATE_ERROR.format("end_with_status", self.state)
             )
         final_status = (
-            self._exc_to_status(status) if isinstance(status, Exception) else status
+            self._exc_to_status(status) if isinstance(status, BaseException) else status
         )
         if self.state == OperationState.ACTIVE_ATTEMPT:
             self.end_attempt_with_status(final_status)
@@ -367,7 +367,7 @@ def cancel(self):
             handler.on_operation_cancelled(self)
 
     @staticmethod
-    def _exc_to_status(exc: Exception) -> StatusCode:
+    def _exc_to_status(exc: BaseException) -> StatusCode:
         """
         Extracts the grpc status code from an exception.
 

tests/system/data/test_metrics_async.py

@@ -1538,8 +1538,8 @@ async def test_sample_row_keys_failure_grpc(
         num_retryable = 3
         for i in range(num_retryable):
             error_injector.push(exc)
-        error_injector.push(RuntimeError)
-        with pytest.raises(RuntimeError):
+        error_injector.push(asyncio.CancelledError)
+        with pytest.raises(asyncio.CancelledError):
             await table.sample_row_keys(retryable_errors=[Aborted])
         # validate counts
         assert len(handler.completed_operations) == 1
@@ -1731,9 +1731,9 @@ async def test_read_modify_write_failure_grpc(
         rule = IncrementRule(family, qualifier, 1)
 
         # trigger an exception
-        exc = RuntimeError("injected")
+        exc = asyncio.CancelledError("injected")
         error_injector.push(exc)
-        with pytest.raises(RuntimeError):
+        with pytest.raises(asyncio.CancelledError):
             await table.read_modify_write_row(row_key, rule)
 
         # validate counts
@@ -1904,9 +1904,9 @@ async def test_check_and_mutate_row_failure_grpc(
         await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier)
 
         # trigger an exception
-        exc = RuntimeError("injected")
+        exc = asyncio.CancelledError("injected")
         error_injector.push(exc)
-        with pytest.raises(RuntimeError):
+        with pytest.raises(asyncio.CancelledError):
             await table.check_and_mutate_row(
                 row_key,
                 predicate=ValueRangeFilter(0, 2),