From 269eaa298f581964819235227673f97dec9ca5e6 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 4 Jun 2018 18:34:32 +0200
Subject: [PATCH] tests: make a fixture out of client

also implement Cancel tests
---
 trezorlib/tests/device_tests/common.py      | 12 ++---
 trezorlib/tests/device_tests/conftest.py    | 52 +++++++++++++++++----
 trezorlib/tests/device_tests/test_cancel.py | 27 +++++++++++
 3 files changed, 73 insertions(+), 18 deletions(-)
 create mode 100644 trezorlib/tests/device_tests/test_cancel.py

diff --git a/trezorlib/tests/device_tests/common.py b/trezorlib/tests/device_tests/common.py
index 5e925b7af..7c31c09fc 100644
--- a/trezorlib/tests/device_tests/common.py
+++ b/trezorlib/tests/device_tests/common.py
@@ -16,28 +16,22 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
-from __future__ import print_function
-
 import os
 
+from . import conftest
+
 from trezorlib import coins
 from trezorlib import tx_api
 from trezorlib.client import TrezorClientDebugLink
-from trezorlib.transport import get_transport
 
 tests_dir = os.path.dirname(os.path.abspath(__file__))
 tx_api.cache_dir = os.path.join(tests_dir, '../txcache')
 
 
-def get_device():
-    path = os.environ.get('TREZOR_PATH')
-    return get_transport(path)
-
-
 class TrezorTest:
 
     def setup_method(self, method):
-        wirelink = get_device()
+        wirelink = conftest.get_device()
         debuglink = wirelink.find_debug()
         self.client = TrezorClientDebugLink(wirelink)
         self.client.set_debuglink(debuglink)
diff --git a/trezorlib/tests/device_tests/conftest.py b/trezorlib/tests/device_tests/conftest.py
index 5f56afe6e..26784a044 100644
--- a/trezorlib/tests/device_tests/conftest.py
+++ b/trezorlib/tests/device_tests/conftest.py
@@ -1,12 +1,19 @@
+import functools
+import os
 import pytest
 
-from . import common
-from trezorlib.client import TrezorClient
-from trezorlib import log
+from trezorlib.transport import get_transport
+from trezorlib.client import TrezorClient, TrezorClientDebugLink
+from trezorlib import log, coins
+
+
+def get_device():
+    path = os.environ.get('TREZOR_PATH')
+    return get_transport(path)
 
 
 def device_version():
-    device = common.get_device()
+    device = get_device()
     if not device:
         raise RuntimeError()
     client = TrezorClient(device)
@@ -16,11 +23,38 @@ def device_version():
         return 1
 
 
-try:
-    TREZOR_VERSION = device_version()
-except:
-    raise
-    TREZOR_VERSION = None
+TREZOR_VERSION = device_version()
+
+
+@pytest.fixture(scope="function")
+def client():
+    wirelink = get_device()
+    debuglink = wirelink.find_debug()
+    client = TrezorClientDebugLink(wirelink)
+    client.set_debuglink(debuglink)
+    client.set_tx_api(coins.tx_api['Bitcoin'])
+    client.wipe_device()
+    client.transport.session_begin()
+
+    yield client
+
+    client.transport.session_end()
+
+
+def setup_client(mnemonic=None, pin='', passphrase=False):
+    if mnemonic is None:
+        mnemonic = ' '.join(['all'] * 12)
+    if pin is True:
+        pin = '1234'
+
+    def client_decorator(function):
+        @functools.wraps(function)
+        def wrapper(client, *args, **kwargs):
+            client.load_device_by_mnemonic(mnemonic=mnemonic, pin=pin, passphrase_protection=passphrase, label='test', language='english')
+            return function(client, *args, **kwargs)
+        return wrapper
+
+    return client_decorator
 
 
 def pytest_configure(config):
diff --git a/trezorlib/tests/device_tests/test_cancel.py b/trezorlib/tests/device_tests/test_cancel.py
new file mode 100644
index 000000000..6d074f4e7
--- /dev/null
+++ b/trezorlib/tests/device_tests/test_cancel.py
@@ -0,0 +1,27 @@
+import pytest
+
+from .conftest import setup_client
+import trezorlib.messages as m
+
+
+@setup_client()
+@pytest.mark.parametrize("message", [
+    m.Ping(message="hello", button_protection=True),
+    m.GetAddress(
+        address_n=[0],
+        coin_name="Bitcoin",
+        script_type=m.InputScriptType.SPENDADDRESS,
+        show_display=True
+    ),
+])
+def test_cancel_ping(client, message):
+    resp = client.call_raw(message)
+    assert isinstance(resp, m.ButtonRequest)
+
+    client.transport.write(m.ButtonAck())
+    client.transport.write(m.Cancel())
+
+    resp = client.transport.read()
+
+    assert isinstance(resp, m.Failure)
+    assert resp.code == m.FailureType.ActionCancelled