diff --git a/tests/common.py b/tests/common.py index f1329b0..8d908bf 100644 --- a/tests/common.py +++ b/tests/common.py @@ -30,6 +30,9 @@ class TrezorTest(unittest.TestCase): def setup_mnemonic_nopin_nopassphrase(self): self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin='', passphrase_protection=False, label='test', language='english') + def setup_mnemonic_pin_nopassphrase(self): + self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin=self.pin4, passphrase_protection=False, label='test', language='english') + def setup_mnemonic_pin_passphrase(self): self.client.load_device_by_mnemonic(mnemonic=self.mnemonic12, pin=self.pin4, passphrase_protection=True, label='test', language='english') diff --git a/tests/test_msg_applysettings.py b/tests/test_msg_applysettings.py index 6c0855b..6bb9bc8 100644 --- a/tests/test_msg_applysettings.py +++ b/tests/test_msg_applysettings.py @@ -5,6 +5,7 @@ import common from trezorlib import messages_pb2 as proto class TestMsgApplysettings(common.TrezorTest): + def test_apply_settings(self): self.setup_mnemonic_pin_passphrase() self.assertEqual(self.client.features.label, 'test') @@ -31,5 +32,37 @@ class TestMsgApplysettings(common.TrezorTest): self.assertEqual(self.client.features.language, 'english') + def test_apply_settings_passphrase(self): + self.setup_mnemonic_pin_nopassphrase() + + self.assertEqual(self.client.features.passphrase_protection, False) + + with self.client: + self.client.set_expected_responses([proto.ButtonRequest(), + proto.PinMatrixRequest(), + proto.Success(), + proto.Features()]) + self.client.apply_settings(use_passphrase=True) + + self.assertEqual(self.client.features.passphrase_protection, True) + + with self.client: + self.client.set_expected_responses([proto.ButtonRequest(), + proto.PinMatrixRequest(), + proto.Success(), + proto.Features()]) + self.client.apply_settings(use_passphrase=False) + + self.assertEqual(self.client.features.passphrase_protection, False) + + with self.client: + self.client.set_expected_responses([proto.ButtonRequest(), + proto.PinMatrixRequest(), + proto.Success(), + proto.Features()]) + self.client.apply_settings(use_passphrase=True) + + self.assertEqual(self.client.features.passphrase_protection, True) + if __name__ == '__main__': unittest.main() diff --git a/trezorlib/client.py b/trezorlib/client.py index 37aa264..2d62db6 100644 --- a/trezorlib/client.py +++ b/trezorlib/client.py @@ -414,12 +414,14 @@ class ProtocolMixin(object): @field('message') @expect(proto.Success) - def apply_settings(self, label=None, language=None): + def apply_settings(self, label=None, language=None, use_passphrase=None): settings = proto.ApplySettings() if label != None: settings.label = label if language: settings.language = language + if use_passphrase != None: + settings.use_passphrase = use_passphrase out = self.call(settings) self.init_device() # Reload Features