Skip to content

Commit

Permalink
Update test_session tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@odeke-em
odeke-em committed Dec 20, 2024
1 parent 57f1570 commit f8e0717
Show file tree
Hide file tree
Showing 4 changed files with 86 additions and 12 deletions.
1 change: 0 additions & 1 deletion google/cloud/spanner_v1/database.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -464,7 +464,6 @@ def _channel_id(self):
return channel_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,
Expand Down
1 change: 1 addition & 0 deletions google/cloud/spanner_v1/transaction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
_merge_query_options,
_metadata_with_prefix,
_metadata_with_leader_aware_routing,
_metadata_with_request_id,
_retry,
_check_rst_stream_error,
)
Expand Down
94 changes: 84 additions & 10 deletions tests/unit/test_session.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,20 @@ def _make_database(name=DATABASE_NAME, database_role=None):
database.log_commit_stats = False
database.database_role = database_role
database._route_to_leader_enabled = True
database.NTH_CLIENT = AtomicCounter()
nth_client_id = AtomicCounter(1)
database.NTH_CLIENT = nth_client_id
next_nth_request = AtomicCounter(1)

def metadata_with_request_id(nth_request, nth_attempt, prior_metadata=[]):
return _metadata_with_request_id(
nth_client_id.value,
1,
next_nth_request.increment(),
nth_attempt,
prior_metadata,
)

database.metadata_with_request_id = metadata_with_request_id
return database

@staticmethod
Expand Down Expand Up @@ -1095,6 +1108,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)

Expand Down Expand Up @@ -1157,10 +1174,25 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
]
* 2,
),
mock.call(
session=self.SESSION_NAME,
options=expected_options,
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}.{database.NTH_CLIENT.value}.1.2.2",
),
],
),
],
)
request = CommitRequest(
session=self.SESSION_NAME,
Expand All @@ -1176,10 +1208,24 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
]
* 2,
),
mock.call(
request=request,
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}.{database.NTH_CLIENT.value}.1.3.2",
),
],
),
],
)

def test_run_in_transaction_w_abort_w_retry_metadata(self):
Expand Down Expand Up @@ -1524,6 +1570,10 @@ def _time(_results=[1, 2, 4, 8]):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
]
Expand All @@ -1543,6 +1593,10 @@ def _time(_results=[1, 2, 4, 8]):
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}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
]
Expand Down Expand Up @@ -1606,6 +1660,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
request = CommitRequest(
Expand All @@ -1620,6 +1678,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
database.logger.info.assert_called_once_with(
Expand Down Expand Up @@ -1676,6 +1738,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
request = CommitRequest(
Expand All @@ -1690,6 +1756,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
database.logger.info.assert_not_called()
Expand Down Expand Up @@ -1753,6 +1823,10 @@ def unit_of_work(txn, *args, **kw):
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}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
request = CommitRequest(
Expand All @@ -1768,7 +1842,7 @@ def unit_of_work(txn, *args, **kw):
("x-goog-spanner-route-to-leader", "true"),
(
"x-goog-spanner-request-id",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.1.1",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
Expand Down Expand Up @@ -1835,7 +1909,7 @@ def unit_of_work(txn, *args, **kw):
("x-goog-spanner-route-to-leader", "true"),
(
"x-goog-spanner-request-id",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.1.1",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.2.1",
),
],
)
Expand All @@ -1852,7 +1926,7 @@ def unit_of_work(txn, *args, **kw):
("x-goog-spanner-route-to-leader", "true"),
(
"x-goog-spanner-request-id",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.1.1",
f"1.{REQ_RAND_PROCESS_ID}.{database.NTH_CLIENT.value}.1.3.1",
),
],
)
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_spanner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1172,7 +1172,7 @@ def test_transaction_should_execute_sql_with_route_to_leader_disabled(self):
(
"x-goog-spanner-request-id",
f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
)
),
],
timeout=TIMEOUT,
retry=RETRY,
Expand Down

0 comments on commit f8e0717

Please sign in to comment.