|
24 | 24 | import uuid |
25 | 25 |
|
26 | 26 | import pytest |
| 27 | +import grpc |
27 | 28 | from google.rpc import code_pb2 |
28 | 29 |
|
29 | 30 | from google.api_core import exceptions |
|
66 | 67 | COUNTERS_TABLE = "counters" |
67 | 68 | COUNTERS_COLUMNS = ("name", "value") |
68 | 69 |
|
| 70 | +_STATUS_CODE_TO_GRPC_STATUS_CODE = { |
| 71 | + member.value[0]: member for member in grpc.StatusCode |
| 72 | +} |
| 73 | + |
69 | 74 |
|
70 | 75 | class Config(object): |
71 | 76 | """Run-time configuration to be modified at set-up. |
@@ -785,9 +790,13 @@ def test_transaction_execute_update_then_insert_commit(self): |
785 | 790 | # [END spanner_test_dml_with_mutation] |
786 | 791 |
|
787 | 792 | @staticmethod |
788 | | - def _check_batch_status(status_code): |
789 | | - if status_code != code_pb2.OK: |
790 | | - raise exceptions.from_grpc_status(status_code, "batch_update failed") |
| 793 | + def _check_batch_status(status_code, expected=code_pb2.OK): |
| 794 | + if status_code != expected: |
| 795 | + grpc_status_code = _STATUS_CODE_TO_GRPC_STATUS_CODE[status_code] |
| 796 | + call = FauxCall(status_code) |
| 797 | + raise exceptions.from_grpc_status( |
| 798 | + grpc_status_code, "batch_update failed", errors=[call] |
| 799 | + ) |
791 | 800 |
|
792 | 801 | def test_transaction_batch_update_success(self): |
793 | 802 | # [START spanner_test_dml_with_mutation] |
@@ -906,7 +915,7 @@ def unit_of_work(transaction): |
906 | 915 | status, row_counts = transaction.batch_update( |
907 | 916 | [insert_statement, update_statement, delete_statement] |
908 | 917 | ) |
909 | | - self.assertEqual(status.code, code_pb2.INVALID_ARGUMENT) |
| 918 | + self._check_batch_status(status.code, code_pb2.INVALID_ARGUMENT) |
910 | 919 | self.assertEqual(len(row_counts), 1) |
911 | 920 | self.assertEqual(row_counts[0], 1) |
912 | 921 |
|
@@ -2190,3 +2199,21 @@ def _handle_abort_unit_of_work(self, transaction): |
2190 | 2199 | def handle_abort(self, database): |
2191 | 2200 | database.run_in_transaction(self._handle_abort_unit_of_work) |
2192 | 2201 | self.handler_done.set() |
| 2202 | + |
| 2203 | + |
| 2204 | +class FauxCall(object): |
| 2205 | + def __init__(self, code, details="FauxCall"): |
| 2206 | + self._code = code |
| 2207 | + self._details = details |
| 2208 | + |
| 2209 | + def initial_metadata(self): |
| 2210 | + return {} |
| 2211 | + |
| 2212 | + def trailing_metadata(self): |
| 2213 | + return {} |
| 2214 | + |
| 2215 | + def code(self): |
| 2216 | + return self._code |
| 2217 | + |
| 2218 | + def details(self): |
| 2219 | + return self._details |
0 commit comments