From 610547cf3211ec0d7910915d80a5189e58248506 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Wed, 10 Dec 2014 15:26:18 +0100
Subject: [PATCH] add button_wait feature for tests

---
 tests/common.py     | 1 +
 trezorlib/client.py | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/tests/common.py b/tests/common.py
index c2b2dd288..f1329b070 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -10,6 +10,7 @@ class TrezorTest(unittest.TestCase):
         self.transport = config.TRANSPORT(*config.TRANSPORT_ARGS, **config.TRANSPORT_KWARGS)
         self.client = TrezorDebugClient(self.transport)
         self.client.set_debuglink(self.debug_transport)
+        # self.client.set_buttonwait(3)
         self.client.set_tx_api(TXAPIBitcoin())
 
         #                     1      2     3    4      5      6      7     8      9    10    11    12
diff --git a/trezorlib/client.py b/trezorlib/client.py
index c7ad6afb2..c6b6e7fb4 100644
--- a/trezorlib/client.py
+++ b/trezorlib/client.py
@@ -209,6 +209,7 @@ class DebugLinkMixin(object):
         super(DebugLinkMixin, self).__init__(*args, **kwargs)
         self.debug = None
         self.in_with_statement = 0
+        self.button_wait = 0
         self.screenshot_id = 0
 
         # Always press Yes and provide correct pin
@@ -228,6 +229,9 @@ class DebugLinkMixin(object):
     def set_debuglink(self, debug_transport):
         self.debug = DebugLink(debug_transport)
 
+    def set_buttonwait(self, secs):
+        self.button_wait = secs
+
     def __enter__(self):
         # For usage in with/expected_responses
         self.in_with_statement += 1
@@ -305,6 +309,9 @@ class DebugLinkMixin(object):
         log("ButtonRequest code: " + get_buttonrequest_value(msg.code))
 
         log("Pressing button " + str(self.button))
+        if self.button_wait:
+            log("Waiting %d seconds " % self.button_wait)
+            time.sleep(self.button_wait)
         self.debug.press_button(self.button)
         return proto.ButtonAck()