|
20 | 20 | import mock |
21 | 21 | import pytest |
22 | 22 |
|
| 23 | +import proto as proto_plus |
| 24 | + |
23 | 25 | from google.cloud.firestore_v1.types import document |
24 | 26 | from google.cloud.firestore_v1.types import firestore |
25 | 27 | from google.cloud.firestore_v1.types import write |
@@ -244,7 +246,10 @@ def callback(keys, applied_changes, read_time): |
244 | 246 | watch = Watch.for_query(query, callback, DocumentSnapshot) |
245 | 247 |
|
246 | 248 | wrapped_responses = [ |
247 | | - firestore.ListenResponse.wrap(proto) for proto in testcase.responses |
| 249 | + firestore.ListenResponse.wrap(proto._pb) |
| 250 | + if isinstance(proto, proto_plus.Message) |
| 251 | + else firestore.ListenResponse.wrap(proto) |
| 252 | + for proto in testcase.responses |
248 | 253 | ] |
249 | 254 | if testcase.is_error: |
250 | 255 | try: |
@@ -336,10 +341,15 @@ def convert_set_option(option): |
336 | 341 | def convert_precondition(precond): |
337 | 342 | from google.cloud.firestore_v1 import Client |
338 | 343 |
|
339 | | - if precond.HasField("exists"): |
| 344 | + if isinstance(precond, proto_plus.Message): |
| 345 | + precond_pb = precond._pb |
| 346 | + else: |
| 347 | + precond_pb = precond |
| 348 | + |
| 349 | + if precond_pb.HasField("exists"): |
340 | 350 | return Client.write_option(exists=precond.exists) |
341 | 351 |
|
342 | | - assert precond.HasField("update_time") |
| 352 | + assert precond_pb.HasField("update_time") |
343 | 353 | return Client.write_option(last_update_time=precond.update_time) |
344 | 354 |
|
345 | 355 |
|
|
0 commit comments