Fix more tests

googleapis · Dec 25, 2024 · 3e47760 · 3e47760
1 parent 27e8fc7
commit 3e47760
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/tests/mockserver_tests/mock_server_test_base.py b/tests/mockserver_tests/mock_server_test_base.py
@@ -134,6 +134,9 @@ def teardown_class(cls):
         if MockServerTestBase.server is not None:
             MockServerTestBase.server.stop(grace=None)
             MockServerTestBase.server = None
+        self.reset()
+
+    def reset(self):
         Client.NTH_CLIENT.reset()
 
     def setup_method(self, *args, **kwargs):

diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py
@@ -1918,7 +1918,7 @@ def test_context_mgr_success(self):
                 ("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",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.1.1.1",
                 ),
             ],
         )
@@ -1969,7 +1969,7 @@ def test_context_mgr_w_commit_stats_success(self):
                 ("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",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.1.1.1",
                 ),
             ],
         )
@@ -2017,7 +2017,7 @@ def test_context_mgr_w_commit_stats_error(self):
                 ("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",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.1.1.1",
                 ),
             ],
         )
@@ -3095,7 +3095,7 @@ def test_context_mgr_success(self):
                 ("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",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.1.1.1",
                 ),
             ],
         )

diff --git a/tests/unit/test_request_id_header.py b/tests/unit/test_request_id_header.py
@@ -50,18 +50,20 @@ def test_snapshot_read(self):
         self.assertTrue(isinstance(requests[0], BatchCreateSessionsRequest))
         self.assertTrue(isinstance(requests[1], ExecuteSqlRequest))
 
+        NTH_CLIENT = self.database._nth_client_id
+        CHANNEL_ID = self.database._channel_id
         # Now ensure monotonicity of the received request-id segments.
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
         want_unary_segments = [
             (
                 "/google.spanner.v1.Spanner/BatchCreateSessions",
-                (1, REQ_RAND_PROCESS_ID, 1, 1, 1, 1),
+                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 1, 1),
             )
         ]
         want_stream_segments = [
             (
                 "/google.spanner.v1.Spanner/ExecuteStreamingSql",
-                (1, REQ_RAND_PROCESS_ID, 1, 1, 2, 1),
+                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 2, 1),
             )
         ]
 
@@ -229,20 +231,22 @@ def test_database_execute_partitioned_dml_request_id(self):
 
         # Now ensure monotonicity of the received request-id segments.
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
+        NTH_CLIENT = self.database._nth_client_id
+        CHANNEL_ID = self.database._channel_id
         want_unary_segments = [
             (
                 "/google.spanner.v1.Spanner/BatchCreateSessions",
-                (1, REQ_RAND_PROCESS_ID, 1, 1, 1, 1),
+                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 1, 1),
             ),
             (
                 "/google.spanner.v1.Spanner/BeginTransaction",
-                (1, REQ_RAND_PROCESS_ID, 1, 1, 2, 1),
+                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 2, 1),
             ),
         ]
         want_stream_segments = [
             (
                 "/google.spanner.v1.Spanner/ExecuteStreamingSql",
-                (1, REQ_RAND_PROCESS_ID, 1, 1, 3, 1),
+                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 3, 1),
             )
         ]