diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index b8827c636..6461df43d 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -209,7 +209,9 @@ def start_attempt(self) -> None: if self.backoff_generator and len(self.completed_attempts) > 0: # find the attempt's backoff by sending attempt number to generator # generator will return the backoff time in seconds, so convert to ms - backoff_ms = self.backoff_generator.send(len(self.completed_attempts) - 1) * 1000 + backoff_ms = ( + self.backoff_generator.send(len(self.completed_attempts) - 1) * 1000 + ) else: backoff_ms = 0 @@ -307,14 +309,14 @@ 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): + status = self._exc_to_status(status) + duration_seconds = time.monotonic() - self.active_attempt.start_time.monotonic new_attempt = CompletedAttemptMetric( start_time=self.active_attempt.start_time.utc, first_response_latency_ms=self.active_attempt.first_response_latency_ms, - duration_ms=(time.monotonic() - self.active_attempt.start_time.monotonic) * 1000, - end_status=self._exc_to_status(status) - if isinstance(status, Exception) - else status, + duration_ms=duration_seconds * 1000, + end_status=status, gfe_latency_ms=self.active_attempt.gfe_latency_ms, application_blocking_time_ms=self.active_attempt.application_blocking_time_ms, backoff_before_attempt_ms=self.active_attempt.backoff_before_attempt_ms, diff --git a/tests/unit/data/_metrics/test_data_model.py b/tests/unit/data/_metrics/test_data_model.py index 5a9dd28df..5e77d5938 100644 --- a/tests/unit/data/_metrics/test_data_model.py +++ b/tests/unit/data/_metrics/test_data_model.py @@ -389,7 +389,9 @@ def test_add_response_metadata_server_timing_header( if metric.active_attempt.gfe_latency_ms is None: assert expected_latency_ms is None else: - assert (metric.active_attempt.gfe_latency_ms - expected_latency_ms) < 0.0001 + assert ( + metric.active_attempt.gfe_latency_ms - expected_latency_ms + ) < 0.0001 # should remain in ACTIVE_ATTEMPT state after completing assert metric.state == State.ACTIVE_ATTEMPT # no errors encountered @@ -708,11 +710,6 @@ def test__handle_error(self): assert logger_mock.warning.call_count == 1 assert logger_mock.warning.call_args[0][0] == expected_message assert len(logger_mock.warning.call_args[0]) == 1 - # otherwise, do nothing - with mock.patch( - "google.cloud.bigtable.data._metrics.data_model.LOGGER", None - ): - type(self._make_one(object()))._handle_error(input_message) @pytest.mark.asyncio async def test_async_context_manager(self):