arno/trezor-firmware
1
0
Fork 0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-01-01 19:10:58 +00:00
Code Issues Releases Wiki Activity

client: clean up exception and decorator code

Browse Source
This commit is contained in:
matejcik 2018-05-11 15:27:39 +02:00
parent a478dac5f7
commit cc7c8ccb59
1 changed files with 12 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
trezorlib/client.py
View File

@ -83,21 +83,14 @@ def get_buttonrequest_value(code):
class CallException(Exception): class CallException(Exception):
def __init__(self, code, message): pass
super(CallException, self).__init__()
self.args = [code, message]
class AssertionException(Exception):
def __init__(self, code, message):
self.args = [code, message]
class PinException(CallException): class PinException(CallException):
pass pass
class field(object): class field:
# Decorator extracts single value from # Decorator extracts single value from
# protobuf object. If the field is not # protobuf object. If the field is not
# present, raises an exception. # present, raises an exception.
@ -105,13 +98,14 @@ class field(object):
self.field = field self.field = field
def __call__(self, f): def __call__(self, f):
@functools.wraps(f)
def wrapped_f(*args, **kwargs): def wrapped_f(*args, **kwargs):
ret = f(*args, **kwargs) ret = f(*args, **kwargs)
return getattr(ret, self.field) return getattr(ret, self.field)
return wrapped_f return wrapped_f
class expect(object): class expect:
# Decorator checks if the method # Decorator checks if the method
# returned one of expected protobuf messages # returned one of expected protobuf messages
# or raises an exception # or raises an exception
@ -119,6 +113,7 @@ class expect(object):
self.expected = expected self.expected = expected
def __call__(self, f): def __call__(self, f):
@functools.wraps(f)
def wrapped_f(*args, **kwargs): def wrapped_f(*args, **kwargs):
ret = f(*args, **kwargs) ret = f(*args, **kwargs)
if not isinstance(ret, self.expected): if not isinstance(ret, self.expected):
@ -130,6 +125,7 @@ class expect(object):
def session(f): def session(f):
# Decorator wraps a BaseClient method # Decorator wraps a BaseClient method
# with session activation / deactivation # with session activation / deactivation
@functools.wraps(f)
def wrapped_f(*args, **kwargs): def wrapped_f(*args, **kwargs):
client = args[0] client = args[0]
client.transport.session_begin() client.transport.session_begin()
@ -388,19 +384,19 @@ class DebugLinkMixin(object):
try: try:
expected = self.expected_responses.pop(0) expected = self.expected_responses.pop(0)
except IndexError: except IndexError:
raise AssertionException(proto.FailureType.UnexpectedMessage, raise AssertionError(proto.FailureType.UnexpectedMessage,
"Got %s, but no message has been expected" % repr(msg)) "Got %s, but no message has been expected" % repr(msg))
if msg.__class__ != expected.__class__: if msg.__class__ != expected.__class__:
raise AssertionException(proto.FailureType.UnexpectedMessage, raise AssertionError(proto.FailureType.UnexpectedMessage,
"Expected %s, got %s" % (repr(expected), repr(msg))) "Expected %s, got %s" % (repr(expected), repr(msg)))
for field, value in expected.__dict__.items(): for field, value in expected.__dict__.items():
if value is None or value == []: if value is None or value == []:
continue continue
if getattr(msg, field) != value: if getattr(msg, field) != value:
raise AssertionException(proto.FailureType.UnexpectedMessage, raise AssertionError(proto.FailureType.UnexpectedMessage,
"Expected %s, got %s" % (repr(expected), repr(msg))) "Expected %s, got %s" % (repr(expected), repr(msg)))
def callback_ButtonRequest(self, msg): def callback_ButtonRequest(self, msg):
self.DEBUG("ButtonRequest code: " + get_buttonrequest_value(msg.code)) self.DEBUG("ButtonRequest code: " + get_buttonrequest_value(msg.code))