Address code review requests

odeke-em · odeke-em · commit 8c27fc6a34d9 · 2025-01-10T04:29:20.000-08:00
diff --git a/tests/system/test_observability_options.py b/tests/system/test_observability_options.py
@@ -302,7 +302,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"]],
@@ -439,18 +439,13 @@ def test_database_partitioned_error():
             codes.ERROR,
             "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
         ),
-        (
-            "CloudSpanner.CreateSession",
-            codes.OK,
-            None,
-        ),
+        ("CloudSpanner.CreateSession", codes.OK, None),
         (
             "CloudSpanner.ExecuteStreamingSql",
             codes.ERROR,
             "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
         ),
     ]
-    print("got_statuses", got_statuses)
     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,13 @@ 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
+        want_span_attributes = dict(TestTransaction.BASE_ATTRIBUTES)
+        want_span_attributes["db.statement"] = DML_QUERY_WITH_PARAM
+        self.assertSpanAttributes(
+            "CloudSpanner.Transaction.execute_update",
+            status=StatusCode.OK,
+            attributes=want_span_attributes,
+        )
 
     def test_execute_update_new_transaction(self):
         self._execute_update_helper()
