1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-16 04:29:08 +00:00
trezor-firmware/tests/device_handler.py
matejcik 5488270bc3 core: improve emulator wrapper
* supply a gzipped image of an empty SD card so that we don't spend 5s at
  every run by formatting it
* add start, stop, restart methods for direct control
* add restart capability to device_handler (a little awkward for now)
2019-10-22 17:06:49 +02:00

85 lines
2.3 KiB
Python

from concurrent.futures import ThreadPoolExecutor
from trezorlib.transport import udp
udp.SOCKET_TIMEOUT = 0.1
class NullUI:
@staticmethod
def button_request(code):
pass
@staticmethod
def get_pin(code=None):
raise NotImplementedError("Should not be used with T1")
@staticmethod
def get_passphrase():
raise NotImplementedError("Should not be used with T1")
class BackgroundDeviceHandler:
_pool = ThreadPoolExecutor()
def __init__(self, client):
self.client = client
self.client.ui = NullUI
self.task = None
def run(self, function, *args, **kwargs):
if self.task is not None:
raise RuntimeError("Wait for previous task first")
self.task = self._pool.submit(function, self.client, *args, **kwargs)
def kill_task(self):
if self.task is not None:
# Force close the client, which should raise an exception in a client
# waiting on IO. Does not work over Bridge, because bridge doesn't have
# a close() method.
while self.client.session_counter > 0:
self.client.close()
try:
self.task.result()
except Exception:
pass
self.task = None
def restart(self, emulator):
# TODO handle actual restart as well
self.kill_task()
emulator.restart()
self.client = emulator.client
self.client.ui = NullUI
def result(self):
if self.task is None:
raise RuntimeError("No task running")
try:
return self.task.result()
finally:
self.task = None
def features(self):
if self.task is not None:
raise RuntimeError("Cannot query features while task is running")
self.client.init_device()
return self.client.features
def debuglink(self):
return self.client.debug
def check_finalize(self):
if self.task is not None:
self.kill_task()
return False
return True
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
finalized_ok = self.check_finalize()
if exc_type is None and not finalized_ok:
raise RuntimeError("Exit while task is unfinished")