fix tests

googleapis · Nov 20, 2024 · f767159 · f767159
1 parent 0c9bff3
commit f767159
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py
@@ -439,7 +439,8 @@ def run_in_transaction(self, func, *args, **kw):
 
                 attempts += 1
 
-                span_attributes = {"transaction.id": txn.id, "attempt": attempts}
+                txn_id = getattr(txn, "_transaction_id", None) or ""
+                span_attributes = {"transaction.id": txn_id, "attempt": attempts}
                 if span:
                     span.add_event("Using Transaction", span_attributes)
 
@@ -448,7 +449,7 @@ def run_in_transaction(self, func, *args, **kw):
                 except Aborted as exc:
                     del self._transaction
                     if span:
-                        delay_seconds = _get_retry_delay(exc, attempts)
+                        delay_seconds = _get_retry_delay(exc.errors[0], attempts)
                         attributes = dict(delay_seconds=delay_seconds)
                         attributes.update(span_attributes)
                         span.add_event("Transaction was aborted, retrying", attributes)
@@ -481,7 +482,7 @@ def run_in_transaction(self, func, *args, **kw):
                 except Aborted as exc:
                     del self._transaction
                     if span:
-                        delay_seconds = _get_retry_delay(exc, attempts)
+                        delay_seconds = _get_retry_delay(exc.errors[0], attempts)
                         attributes = dict(delay_seconds=delay_seconds)
                         attributes.update(span_attributes)
                         span.add_event(

diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py
@@ -1194,11 +1194,8 @@ def unit_of_work(transaction):
     ]
     assert [span.name for span in span_list] == expected_span_names
     for span in span_list[2:-2]:
-        assert (
-            span.get_span_context().trace_id
-            == span_list[-1].get_span_context().trace_id
-        )
-        assert span.parent.span_id == span_list[-2].get_span_context().span_id
+        assert span.context.trace_id == span_list[-2].context.trace_id
+        assert span.parent.span_id == span_list[-2].context.span_id
 
 
 def test_execute_partitioned_dml(