1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-02-21 20:12:02 +00:00

tests: only allow input_flow in a with-block

This commit is contained in:
matejcik 2019-09-18 16:13:11 +02:00
parent eca0a719db
commit 4eb98c254a
4 changed files with 46 additions and 42 deletions

View File

@ -237,6 +237,9 @@ class TrezorClientDebugLink(TrezorClient):
self.ui.input_flow = None self.ui.input_flow = None
return return
if not self.in_with_statement:
raise RuntimeError("Must be called inside 'with' statement")
if callable(input_flow): if callable(input_flow):
input_flow = input_flow() input_flow = input_flow()
if not hasattr(input_flow, "send"): if not hasattr(input_flow, "send"):
@ -252,6 +255,9 @@ class TrezorClientDebugLink(TrezorClient):
def __exit__(self, _type, value, traceback): def __exit__(self, _type, value, traceback):
self.in_with_statement -= 1 self.in_with_statement -= 1
# Clear input flow.
self.set_input_flow(None)
if _type is not None: if _type is not None:
# Another exception raised # Another exception raised
return False return False
@ -268,6 +274,7 @@ class TrezorClientDebugLink(TrezorClient):
# Cleanup # Cleanup
self.expected_responses = None self.expected_responses = None
self.current_response = None self.current_response = None
return False return False
def set_expected_responses(self, expected): def set_expected_responses(self, expected):

View File

@ -194,24 +194,26 @@ class TestMsgTezosSignTx:
debug.press_yes() debug.press_yes()
def test_tezos_sign_tx_proposal(self, client): def test_tezos_sign_tx_proposal(self, client):
client.set_input_flow(self.input_flow(client.debug, num_pages=1)) with client:
resp = tezos.sign_tx( client.set_input_flow(self.input_flow(client.debug, num_pages=1))
client, resp = tezos.sign_tx(
TEZOS_PATH_10, client,
dict_to_proto( TEZOS_PATH_10,
messages.TezosSignTx, dict_to_proto(
{ messages.TezosSignTx,
"branch": "dee04042c0832d68a43699b2001c0a38065436eb05e578071a763e1972d0bc81", {
"proposal": { "branch": "dee04042c0832d68a43699b2001c0a38065436eb05e578071a763e1972d0bc81",
"source": "005f450441f41ee11eee78a31d1e1e55627c783bd6", "proposal": {
"period": 17, "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
"proposals": [ "period": 17,
"dfa974df171c2dad9a9b8f25d99af41fd9702ce5d04521d2f9943c84d88aa572" "proposals": [
], "dfa974df171c2dad9a9b8f25d99af41fd9702ce5d04521d2f9943c84d88aa572"
],
},
}, },
}, ),
), )
)
assert ( assert (
resp.signature resp.signature
== "edsigtfY16R32k2WVMYfFr7ymnro4ib5zMckk28vsuViYNN77DJAvCJLRNArd9L531pUCxT4YdcvCvBym5dhcZ1rknEVm6yZ8bB" == "edsigtfY16R32k2WVMYfFr7ymnro4ib5zMckk28vsuViYNN77DJAvCJLRNArd9L531pUCxT4YdcvCvBym5dhcZ1rknEVm6yZ8bB"
@ -225,25 +227,27 @@ class TestMsgTezosSignTx:
) )
def test_tezos_sign_tx_multiple_proposals(self, client): def test_tezos_sign_tx_multiple_proposals(self, client):
client.set_input_flow(self.input_flow(client.debug, num_pages=2)) with client:
resp = tezos.sign_tx( client.set_input_flow(self.input_flow(client.debug, num_pages=2))
client, resp = tezos.sign_tx(
TEZOS_PATH_10, client,
dict_to_proto( TEZOS_PATH_10,
messages.TezosSignTx, dict_to_proto(
{ messages.TezosSignTx,
"branch": "7e0be36a90c663c73c60da3889ffefff1383fb65cc29f0639f173d8f95a52df7", {
"proposal": { "branch": "7e0be36a90c663c73c60da3889ffefff1383fb65cc29f0639f173d8f95a52df7",
"source": "005f450441f41ee11eee78a31d1e1e55627c783bd6", "proposal": {
"period": 17, "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
"proposals": [ "period": 17,
"2a6ff28ab4d0ccb18f7129aaaf9a4b8027d794f2562849665fdb6999db2a4e57", "proposals": [
"47cd60c09ab8437cc9fe19add494dce1b9844100f660f02ce77510a0c66d2762", "2a6ff28ab4d0ccb18f7129aaaf9a4b8027d794f2562849665fdb6999db2a4e57",
], "47cd60c09ab8437cc9fe19add494dce1b9844100f660f02ce77510a0c66d2762",
],
},
}, },
}, ),
), )
)
assert ( assert (
resp.signature resp.signature
== "edsigu6GAjhiWAQ64ctWTGEDYAZ16tYzLgzWzqc4CUyixK4FGRE8YUBVzFaVJ2fUCexZjZLMLdiNZGcUdzeL1bQhZ2h5oLrh7pA" == "edsigu6GAjhiWAQ64ctWTGEDYAZ16tYzLgzWzqc4CUyixK4FGRE8YUBVzFaVJ2fUCexZjZLMLdiNZGcUdzeL1bQhZ2h5oLrh7pA"

View File

@ -179,7 +179,6 @@ def reset(client, strength=128):
language="english", language="english",
backup_type=BackupType.Slip39_Advanced, backup_type=BackupType.Slip39_Advanced,
) )
client.set_input_flow(None)
# Check if device is properly initialized # Check if device is properly initialized
assert client.features.initialized is True assert client.features.initialized is True
@ -203,8 +202,6 @@ def recover(client, shares):
client.set_input_flow(input_flow) client.set_input_flow(input_flow)
ret = device.recover(client, pin_protection=False, label="label") ret = device.recover(client, pin_protection=False, label="label")
client.set_input_flow(None)
# Workflow successfully ended # Workflow successfully ended
assert ret == messages.Success(message="Device recovered") assert ret == messages.Success(message="Device recovered")
assert client.features.pin_protection is False assert client.features.pin_protection is False

View File

@ -115,8 +115,6 @@ def reset(client, strength=128):
backup_type=BackupType.Slip39_Basic, backup_type=BackupType.Slip39_Basic,
) )
client.set_input_flow(None)
# Check if device is properly initialized # Check if device is properly initialized
assert client.features.initialized is True assert client.features.initialized is True
assert client.features.needs_backup is False assert client.features.needs_backup is False
@ -139,8 +137,6 @@ def recover(client, shares):
client.set_input_flow(input_flow) client.set_input_flow(input_flow)
ret = device.recover(client, pin_protection=False, label="label") ret = device.recover(client, pin_protection=False, label="label")
client.set_input_flow(None)
# Workflow successfully ended # Workflow successfully ended
assert ret == messages.Success(message="Device recovered") assert ret == messages.Success(message="Device recovered")
assert client.features.pin_protection is False assert client.features.pin_protection is False