added sync tests

Author: daniel-sanche

google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py

@@ -15,7 +15,9 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
+import time
 from functools import wraps
+from grpc import RpcError
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -42,7 +44,10 @@ def wrapper(self, continuation, client_call_details, request):
         )
         operation: "ActiveOperationMetric" = self.operation_map.get(key)
         if operation:
-            if operation.state != OperationState.ACTIVE_ATTEMPT:
+            if (
+                operation.state == OperationState.CREATED
+                or operation.state == OperationState.BETWEEN_ATTEMPTS
+            ):
                 operation.start_attempt()
             return func(self, operation, continuation, client_call_details, request)
         else:
@@ -76,7 +81,8 @@ def register_operation(self, operation):
         operation.handlers.append(self)
 
     def on_operation_complete(self, op):
-        del self.operation_map[op.uuid]
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
 
     def on_operation_cancelled(self, op):
         self.on_operation_complete(op)
@@ -86,36 +92,70 @@ def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
         encountered_exc: Exception | None = None
-        call = None
+        metadata = None
         try:
             call = continuation(client_call_details, request)
+            metadata = (call.trailing_metadata() or []) + (
+                call.initial_metadata() or []
+            )
             return call
+        except RpcError as rpc_error:
+            try:
+                metadata = (rpc_error.trailing_metadata() or []) + (
+                    rpc_error.initial_metadata() or []
+                )
+            except Exception:
+                pass
+            encountered_exc = rpc_error
+            raise rpc_error
         except Exception as e:
             encountered_exc = e
             raise
         finally:
-            if call is not None:
-                metadata = call.trailing_metadata() + call.initial_metadata()
+            if metadata is not None:
                 operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    operation.end_attempt_with_status(encountered_exc)
+            if encountered_exc is not None:
+                operation.end_attempt_with_status(encountered_exc)
 
     @_with_operation_from_metadata
     def intercept_unary_stream(
         self, operation, continuation, client_call_details, request
     ):
         def response_wrapper(call):
+            has_first_response = operation.first_response_latency is not None
             encountered_exc = None
             try:
                 for response in call:
+                    if not has_first_response:
+                        operation.first_response_latency_ns = (
+                            time.monotonic_ns() - operation.start_time_ns
+                        )
+                        has_first_response = True
                     yield response
             except Exception as e:
                 encountered_exc = e
                 raise
             finally:
-                metadata = call.trailing_metadata() + call.initial_metadata()
-                operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    operation.end_attempt_with_status(encountered_exc)
+                if call is not None:
+                    metadata = (call.trailing_metadata() or []) + (
+                        call.initial_metadata() or []
+                    )
+                    operation.add_response_metadata(metadata)
+                    if encountered_exc is not None:
+                        operation.end_attempt_with_status(encountered_exc)
 
-        return response_wrapper(continuation(client_call_details, request))
+        try:
+            return response_wrapper(continuation(client_call_details, request))
+        except RpcError as rpc_error:
+            try:
+                metadata = (rpc_error.trailing_metadata() or []) + (
+                    rpc_error.initial_metadata() or []
+                )
+                operation.add_response_metadata(metadata)
+            except Exception:
+                pass
+            operation.end_attempt_with_status(rpc_error)
+            raise rpc_error
+        except Exception as e:
+            operation.end_attempt_with_status(e)
+            raise

tests/unit/data/_async/test_metrics_interceptor.py

@@ -25,6 +25,11 @@
 except ImportError:  # pragma: NO COVER
     import mock  # type: ignore
 
+if CrossSync.is_async:
+    from google.cloud.bigtable.data._async.metrics_interceptor import AsyncBigtableMetricsInterceptor
+else:
+    from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import BigtableMetricsInterceptor
+
 
 __CROSS_SYNC_OUTPUT__ = "tests.unit.data._sync_autogen.test_metrics_interceptor"
 
@@ -53,11 +58,9 @@ async def __anext__(self):
 @CrossSync.convert_class(sync_name="TestMetricsInterceptor")
 class TestMetricsInterceptorAsync:
     @staticmethod
-    @CrossSync.convert
+    @CrossSync.convert(replace_symbols={"AsyncBigtableMetricsInterceptor": "BigtableMetricsInterceptor"})
     def _get_target_class():
-        from google.cloud.bigtable.data._async import metrics_interceptor
-
-        return metrics_interceptor.AsyncBigtableMetricsInterceptor
+        return AsyncBigtableMetricsInterceptor
 
     def _make_one(self, *args, **kwargs):
         return self._get_target_class()(*args, **kwargs)