Address code review requests

googleapis · Jan 10, 2025 · 68cae9d · 68cae9d
1 parent c953e5b
commit 68cae9d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 87 deletions.
diff --git a/tests/system/test_observability_options.py b/tests/system/test_observability_options.py
@@ -298,7 +298,7 @@ def test_transaction_update_implicit_begin_nested_inside_commit():
     LABELS = {"test": "true"}
 
     def tx_update(txn):
-        txn.update(
+        txn.insert(
             "Singers",
             columns=["SingerId", "FirstName"],
             values=[["1", "Bryan"], ["2", "Slash"]],
@@ -403,60 +403,19 @@ def tx_update(txn):
         ("Waiting for a session to become available", {"kind": "BurstyPool"}),
         ("No sessions available in pool. Creating session", {"kind": "BurstyPool"}),
         ("Creating Session", {}),
-        (
-            "exception",
-            {
-                "exception.type": "google.api_core.exceptions.NotFound",
-                "exception.message": "404 Table Singers: Row {Int64(1)} not found.",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
-        (
-            "Transaction.commit failed due to GoogleAPICallError, not retrying",
-            {"attempt": 1},
-        ),
-        (
-            "exception",
-            {
-                "exception.type": "google.api_core.exceptions.NotFound",
-                "exception.message": "404 Table Singers: Row {Int64(1)} not found.",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
         ("Starting Commit", {}),
-        (
-            "exception",
-            {
-                "exception.type": "google.api_core.exceptions.NotFound",
-                "exception.message": "404 Table Singers: Row {Int64(1)} not found.",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
+        ("Commit Done", {}),
     ]
+    print("got_events", got_events)
     assert got_events == want_events
 
     # Check for the statues.
     codes = StatusCode
     want_statuses = [
-        (
-            "CloudSpanner.Database.run_in_transaction",
-            codes.ERROR,
-            "NotFound: 404 Table Singers: Row {Int64(1)} not found.",
-        ),
+        ("CloudSpanner.Database.run_in_transaction", codes.OK, None),
         ("CloudSpanner.CreateSession", codes.OK, None),
-        (
-            "CloudSpanner.Session.run_in_transaction",
-            codes.ERROR,
-            "NotFound: 404 Table Singers: Row {Int64(1)} not found.",
-        ),
-        (
-            "CloudSpanner.Transaction.commit",
-            codes.ERROR,
-            "NotFound: 404 Table Singers: Row {Int64(1)} not found.",
-        ),
+        ("CloudSpanner.Session.run_in_transaction", codes.OK, None),
+        ("CloudSpanner.Transaction.commit", codes.OK, None),
         ("CloudSpanner.Transaction.begin", codes.OK, None),
     ]
     assert got_statuses == want_statuses

diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
 
 import mock
 from google.api_core import gapic_v1
@@ -22,10 +23,6 @@
 from google.cloud.spanner_v1.param_types import INT64
 from google.api_core.retry import Retry
 from google.protobuf.field_mask_pb2 import FieldMask
-from tests._helpers import (
-    HAS_OPENTELEMETRY_INSTALLED,
-    OpenTelemetryBase,
-)
 
 from google.cloud.spanner_v1 import RequestOptions, DirectedReadOptions
 
@@ -64,7 +61,7 @@ class _CredentialsWithScopes(
     return mock.Mock(spec=_CredentialsWithScopes)
 
 
-class _BaseTest(OpenTelemetryBase):
+class _BaseTest(unittest.TestCase):
     PROJECT_ID = "project-id"
     PARENT = "projects/" + PROJECT_ID
     INSTANCE_ID = "instance-id"
@@ -1435,18 +1432,6 @@ def test_run_in_transaction_w_args(self):
         self.assertEqual(committed, NOW)
         self.assertEqual(session._retried, (_unit_of_work, (SINCE,), {"until": UNTIL}))
 
-        if not HAS_OPENTELEMETRY_INSTALLED:
-            pass
-
-        span_list = self.get_finished_spans()
-        got_span_names = [span.name for span in span_list]
-        want_span_names = ["CloudSpanner.Database.run_in_transaction"]
-        assert got_span_names == want_span_names
-
-        got_span_events_statuses = self.finished_spans_events_statuses()
-        want_span_events_statuses = []
-        assert got_span_events_statuses == want_span_events_statuses
-
     def test_run_in_transaction_nested(self):
         from datetime import datetime
 

diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py
@@ -229,6 +229,7 @@ def test_rollback_not_begun(self):
 
         # Since there was no transaction to be rolled back, rollback rpc is not called.
         api.rollback.assert_not_called()
+
         self.assertNoSpans()
 
     def test_rollback_already_committed(self):
@@ -298,17 +299,10 @@ def test_rollback_ok(self):
             ],
         )
 
-        if not HAS_OPENTELEMETRY_INSTALLED:
-            pass
-
-        span_list = self.get_finished_spans()
-        got_span_names = [span.name for span in span_list]
-        want_span_names = ["CloudSpanner.Transaction.rollback"]
-        assert got_span_names == want_span_names
-
-        got_span_events_statuses = self.finished_spans_events_statuses()
-        want_span_events_statuses = []
-        assert got_span_events_statuses == want_span_events_statuses
+        self.assertSpanAttributes(
+            "CloudSpanner.Transaction.rollback",
+            attributes=TestTransaction.BASE_ATTRIBUTES,
+        )
 
     def test_commit_not_begun(self):
         session = _Session()
@@ -665,18 +659,11 @@ def _execute_update_helper(
         )
 
         self.assertEqual(transaction._execute_sql_count, count + 1)
-
-        if not HAS_OPENTELEMETRY_INSTALLED:
-            pass
-
-        span_list = self.get_finished_spans()
-        got_span_names = [span.name for span in span_list]
-        want_span_names = ["CloudSpanner.Transaction.execute_update"]
-        assert got_span_names == want_span_names
-
-        got_span_events_statuses = self.finished_spans_events_statuses()
-        want_span_events_statuses = []
-        assert got_span_events_statuses == want_span_events_statuses
+        self.assertSpanAttributes(
+            "CloudSpanner.Transaction.execute_update",
+            status=StatusCode.ERROR,
+            attributes=TestTransaction.BASE_ATTRIBUTES,
+        )
 
     def test_execute_update_new_transaction(self):
         self._execute_update_helper()