From 7ebbccdbba3eb5d833c94d5e82a0348c80ada8a3 Mon Sep 17 00:00:00 2001 From: matejcik Date: Fri, 24 Apr 2020 12:58:17 +0200 Subject: [PATCH] python/debuglink: support conditional expected_responses --- python/src/trezorlib/debuglink.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index d7dd4aac8..68c78b6c3 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -384,10 +384,31 @@ class TrezorClientDebugLink(TrezorClient): must exactly match the received field value. If a given field is None (or unspecified) in the expected response, the received field value is not checked. + + Each expected response can also be a tuple (bool, message). In that case, the + expected response is only evaluated if the first field is True. + This is useful for differentiating sequences between Trezor models: + + >>> trezor_one = client.features.model == "1" + >>> client.set_expected_responses([ + >>> messages.ButtonRequest(code=ConfirmOutput), + >>> (trezor_one, messages.ButtonRequest(code=ConfirmOutput)), + >>> messages.Success(), + >>> ]) """ if not self.in_with_statement: raise RuntimeError("Must be called inside 'with' statement") - self.expected_responses = expected + + # make sure all items are (bool, message) tuples + expected_with_validity = [ + e if isinstance(e, tuple) else (True, e) for e in expected + ] + + # only apply those items that are (True, message) + self.expected_responses = [ + expected for valid, expected in expected_with_validity if valid + ] + self.current_response = 0 def use_pin_sequence(self, pins):