diff --git a/tests/test_entropy.py b/tests/test_entropy.py index 4183d63..831a70b 100644 --- a/tests/test_entropy.py +++ b/tests/test_entropy.py @@ -2,6 +2,9 @@ import unittest import common import math +import trezorlib.messages_pb2 as proto +import trezorlib.types_pb2 as proto_types + def entropy(data): counts = {} for c in data: @@ -19,6 +22,7 @@ class TestEntropy(common.TrezorTest): def test_entropy(self): for l in [0, 1, 2, 3, 4, 5, 8, 9, 16, 17, 32, 33, 64, 65, 128, 129, 256, 257, 512, 513, 1024]: + self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_Other), proto.Entropy()]) ent = self.client.get_entropy(l) self.assertTrue(len(ent) >= l) print 'entropy = ', entropy(ent) diff --git a/tests/test_pings.py b/tests/test_pings.py index 2dc72cc..0d6bf7f 100644 --- a/tests/test_pings.py +++ b/tests/test_pings.py @@ -3,33 +3,52 @@ import unittest import common from trezorlib import messages_pb2 as proto -from trezorlib import types_pb2 as types +from trezorlib import types_pb2 as proto_types class TestPing(common.TrezorTest): def test_ping(self): - + self.client.wipe_device() + self.client.load_device_by_mnemonic( + mnemonic=self.mnemonic1, + pin=self.pin1, + passphrase_protection=True, + label='test', + language='english', + ) + + self.client.set_expected_responses([proto.Success()]) res = self.client.ping('random data') self.assertEqual(res, 'random data') + self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_Other),proto.Success()]) res = self.client.ping('random data', button_protection=True) self.assertEqual(res, 'random data') + self.client.set_expected_responses([proto.PinMatrixRequest(),proto.Success()]) res = self.client.ping('random data', pin_protection=True) self.assertEqual(res, 'random data') + self.client.set_expected_responses([proto.PassphraseRequest(),proto.Success()]) res = self.client.ping('random data', passphrase_protection=True) self.assertEqual(res, 'random data') - res = self.client.ping('random data', button_protection=True, pin_protection=True) - self.assertEqual(res, 'random data') + def test_ping_caching(self): + self.client.wipe_device() + self.client.load_device_by_mnemonic( + mnemonic=self.mnemonic1, + pin=self.pin1, + passphrase_protection=True, + label='test', + language='english', + ) - res = self.client.ping('random data', button_protection=True, passphrase_protection=True) - self.assertEqual(res, 'random data') - - res = self.client.ping('random data', pin_protection=True, passphrase_protection=True) + self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_Other),proto.PinMatrixRequest(),proto.PassphraseRequest(),proto.Success()]) + res = self.client.ping('random data', button_protection=True, pin_protection=True, passphrase_protection=True) self.assertEqual(res, 'random data') + # pin and passphrase are cached + self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_Other)]) res = self.client.ping('random data', button_protection=True, pin_protection=True, passphrase_protection=True) self.assertEqual(res, 'random data')