Plumb in more tests

Author: odeke-em

tests/unit/test_snapshot.py

@@ -25,6 +25,11 @@
 )
 from google.cloud.spanner_v1.param_types import INT64
 from google.api_core.retry import Retry
+from google.cloud.spanner_v1._helpers import (
+    AtomicCounter,
+    _metadata_with_request_id,
+)
+from google.cloud.spanner_v1.request_id_header import REQ_RAND_PROCESS_ID
 
 TABLE_NAME = "citizens"
 COLUMNS = ["email", "first_name", "last_name", "age"]
@@ -278,7 +283,7 @@ def test_iteration_w_raw_raising_retryable_internal_error(self):
             fail_after=True,
             error=InternalServerError(
                 "Received unexpected EOS on DATA frame from server"
-            )
+            ),
         )
         after = _MockIterator(*LAST)
         request = mock.Mock(test="test", spec=["test", "resume_token"])
@@ -450,7 +455,7 @@ def test_iteration_w_raw_raising_retryable_internal_error_after_token(self):
             fail_after=True,
             error=InternalServerError(
                 "Received unexpected EOS on DATA frame from server"
-            )
+            ),
         )
         after = _MockIterator(*SECOND)
         request = mock.Mock(test="test", spec=["test", "resume_token"])
@@ -767,7 +772,13 @@ def _read_helper(
         )
         api.streaming_read.assert_called_once_with(
             request=expected_request,
-            metadata=[("google-cloud-resource-prefix", database.name)],
+            metadata=[
+                ("google-cloud-resource-prefix", database.name),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
+            ],
             retry=retry,
             timeout=timeout,
         )
@@ -1016,7 +1027,13 @@ def _execute_sql_helper(
         )
         api.execute_streaming_sql.assert_called_once_with(
             request=expected_request,
-            metadata=[("google-cloud-resource-prefix", database.name)],
+            metadata=[
+                ("google-cloud-resource-prefix", database.name),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
+            ],
             timeout=timeout,
             retry=retry,
         )
@@ -1189,6 +1206,10 @@ def _partition_read_helper(
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
             retry=retry,
             timeout=timeout,
@@ -1363,6 +1384,10 @@ def _partition_query_helper(
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
             retry=retry,
             timeout=timeout,
@@ -1751,7 +1776,13 @@ def test_begin_ok_exact_staleness(self):
         api.begin_transaction.assert_called_once_with(
             session=session.name,
             options=expected_txn_options,
-            metadata=[("google-cloud-resource-prefix", database.name)],
+            metadata=[
+                ("google-cloud-resource-prefix", database.name),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
+            ],
         )
 
         self.assertSpanAttributes(
@@ -1787,7 +1818,13 @@ def test_begin_ok_exact_strong(self):
         api.begin_transaction.assert_called_once_with(
             session=session.name,
             options=expected_txn_options,
-            metadata=[("google-cloud-resource-prefix", database.name)],
+            metadata=[
+                ("google-cloud-resource-prefix", database.name),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
+            ],
         )
 
         self.assertSpanAttributes(
@@ -1798,10 +1835,18 @@ def test_begin_ok_exact_strong(self):
 
 
 class _Client(object):
+    NTH_CLIENT = AtomicCounter()
+
     def __init__(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        self._nth_client_id = _Client.NTH_CLIENT.increment()
+        self._nth_request = AtomicCounter()
+
+    @property
+    def _next_nth_request(self):
+        return self._nth_request.increment()
 
 
 class _Instance(object):
@@ -1816,6 +1861,28 @@ def __init__(self, directed_read_options=None):
         self._route_to_leader_enabled = True
         self._directed_read_options = directed_read_options
 
+    @property
+    def _next_nth_request(self):
+        return self._instance._client._next_nth_request
+
+    @property
+    def _nth_client_id(self):
+        return self._instance._client._nth_client_id
+
+    def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
+        client_id = self._nth_client_id
+        return _metadata_with_request_id(
+            self._nth_client_id,
+            self._channel_id,
+            nth_request,
+            nth_attempt,
+            prior_metadata,
+        )
+
+    @property
+    def _channel_id(self):
+        return 1
+
 
 class _Session(object):
     def __init__(self, database=None, name=TestSnapshot.SESSION_NAME):