Implements 'with' + set_expected_responses

tests/test.py

@@ -21,6 +21,8 @@ x SimpleSignTx
 FirmwareErase
 FirmwareUpload
 
+protection levels
+
 - zrejme v sucinnosti s inymi testami
   x ButtonRequest/ButtonAck workflow
   x PinMatrixRequest/PinMatrixAck workflow

tests/test_bip32_speed.py

@@ -3,7 +3,7 @@ import common
 import time
 from trezorlib import tools
 
-class TestAddresses(common.TrezorTest):
+class TestBip32Speed(common.TrezorTest):
     def test_public_ckd(self):
         self.setup_mnemonic_nopin_nopassphrase()
 
@@ -13,7 +13,7 @@ class TestAddresses(common.TrezorTest):
             start = time.time()
             self.client.get_address('Bitcoin', range(depth))
             delay = time.time() - start
-            expected = (depth + 1) * 0.25
+            expected = (depth + 1) * 0.26
             print "DEPTH", depth, "EXPECTED DELAY", expected, "REAL DELAY", delay
             self.assertLessEqual(delay, expected)
 
@@ -26,7 +26,7 @@ class TestAddresses(common.TrezorTest):
             start = time.time()
             self.client.get_address('Bitcoin', range(-depth, 0))
             delay = time.time() - start
-            expected = (depth + 1) * 0.25
+            expected = (depth + 1) * 0.26
             print "DEPTH", depth, "EXPECTED DELAY", expected, "REAL DELAY", delay
             self.assertLessEqual(delay, expected)
 

tests/test_msg_getaddress.py

@@ -3,7 +3,7 @@ import common
 import trezorlib.ckd_public as bip32
 from trezorlib import tools
 
-class TestAddresses(common.TrezorTest):
+class TestMsgGetaddress(common.TrezorTest):
 
     def test_btc(self):
         self.setup_mnemonic_nopin_nopassphrase()

tests/test_msg_getentropy.py

@@ -18,14 +18,15 @@ def entropy(data):
         e -= p * math.log(p, 256)
     return e
 
-class TestEntropy(common.TrezorTest):
+class TestMsgGetentropy(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)
+            with self.client:
+                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)
 
 if __name__ == '__main__':
     unittest.main()

tests/test_msg_ping.py

@@ -10,33 +10,39 @@ class TestPing(common.TrezorTest):
     def test_ping(self):
         self.setup_mnemonic_pin_passphrase()
 
-        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')
+        with self.client:
+            self.client.set_expected_responses([proto.Success()])
+            res = self.client.ping('random data')
+            self.assertEqual(res, 'random data')
+
+        with self.client:
+            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')
+
+        with self.client:
+            self.client.set_expected_responses([proto.PinMatrixRequest(), proto.Success()])
+            res = self.client.ping('random data', pin_protection=True)
+            self.assertEqual(res, 'random data')
+
+        with self.client:
+            self.client.set_expected_responses([proto.PassphraseRequest(), proto.Success()])
+            res = self.client.ping('random data', passphrase_protection=True)
+            self.assertEqual(res, 'random data')
 
     def test_ping_caching(self):
         self.setup_mnemonic_pin_passphrase()
 
-        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')
+        with self.client:
+            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),proto.Success()])
-        res = self.client.ping('random data', button_protection=True, pin_protection=True, passphrase_protection=True)
-        self.assertEqual(res, 'random data')
+        with self.client:
+            # pin and passphrase are cached
+            self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_Other), proto.Success()])
+            res = self.client.ping('random data', button_protection=True, pin_protection=True, passphrase_protection=True)
+            self.assertEqual(res, 'random data')
 
 if __name__ == '__main__':
     unittest.main()

tests/test_protect_call.py

@@ -23,49 +23,59 @@ class TestProtectCall(common.TrezorTest):
         # This is low-level test of set_expected_responses()
         # feature of debugging client
 
-        # Scenario 1 - Received unexpected message
-        self.client.set_expected_responses([])
-        self.assertRaises(CallException, self._some_protected_call, True, True, True)
-
-        # Scenario 2 - Received other than expected message
-        self.client.set_expected_responses([proto.Success()])
-        self.assertRaises(CallException, self._some_protected_call, True, True, True)
-
-        # Scenario 3 - Not received expected message
-        self.client.set_expected_responses([proto.ButtonRequest(),
-                                            proto.Success(),
-                                            proto.Success()])  # This is expected, but not received
-        self.assertRaises(Exception, self._some_protected_call, True, False, False)
-
-        # Scenario 4 - Received what expected
-        self.client.set_expected_responses([proto.ButtonRequest(),
-                                            proto.PinMatrixRequest(),
-                                            proto.PassphraseRequest(),
-                                            proto.Success(message='random data')])
-        self._some_protected_call(True, True, True)
-
-        # Scenario 5 - Failed message by field filter
-        self.client.set_expected_responses([proto.ButtonRequest(),
-                                            proto.PinMatrixRequest(),
-                                            proto.Success(message='wrong data')])
-        self.assertRaises(CallException, self._some_protected_call, True, True, True)
+        with self.client:
+            # Scenario 1 - Received unexpected message
+            self.client.set_expected_responses([])
+            self.assertRaises(CallException, self._some_protected_call, True, True, True)
+
+        with self.client:
+            # Scenario 2 - Received other than expected message
+            self.client.set_expected_responses([proto.Success()])
+            self.assertRaises(CallException, self._some_protected_call, True, True, True)
+
+        def scenario3():
+            with self.client:
+                # Scenario 3 - Not received expected message
+                self.client.set_expected_responses([proto.ButtonRequest(),
+                                                    proto.Success(),
+                                                    proto.Success()])  # This is expected, but not received
+                self._some_protected_call(True, False, False)
+        self.assertRaises(Exception, scenario3)
+
+        with self.client:
+            # Scenario 4 - Received what expected
+            self.client.set_expected_responses([proto.ButtonRequest(),
+                                                proto.PinMatrixRequest(),
+                                                proto.PassphraseRequest(),
+                                                proto.Success(message='random data')])
+            self._some_protected_call(True, True, True)
+
+        def scenario5():
+            with self.client:
+                # Scenario 5 - Failed message by field filter
+                self.client.set_expected_responses([proto.ButtonRequest(),
+                                                    proto.Success(message='wrong data')])
+                self._some_protected_call(True, True, True)
+        self.assertRaises(CallException, scenario5)
 
     def test_no_protection(self):
         self.setup_mnemonic_nopin_nopassphrase()
 
-        self.assertEqual(self.client.debug.read_pin()[0], '')
-        self.client.set_expected_responses([proto.Success()])
-        self._some_protected_call(False, True, True)
+        with self.client:
+            self.assertEqual(self.client.debug.read_pin()[0], '')
+            self.client.set_expected_responses([proto.Success()])
+            self._some_protected_call(False, True, True)
 
     def test_pin(self):
         self.setup_mnemonic_pin_passphrase()
 
-        self.assertEqual(self.client.debug.read_pin()[0], self.pin4)
-        self.client.setup_debuglink(button=True, pin_correct=True)
-        self.client.set_expected_responses([proto.ButtonRequest(),
-                                            proto.PinMatrixRequest(),
-                                            proto.Success()])
-        self._some_protected_call(True, True, False)
+        with self.client:
+            self.assertEqual(self.client.debug.read_pin()[0], self.pin4)
+            self.client.setup_debuglink(button=True, pin_correct=True)
+            self.client.set_expected_responses([proto.ButtonRequest(),
+                                                proto.PinMatrixRequest(),
+                                                proto.Success()])
+            self._some_protected_call(True, True, False)
 
     def test_incorrect_pin(self):
         self.setup_mnemonic_pin_passphrase()
@@ -94,6 +104,7 @@ class TestProtectCall(common.TrezorTest):
             start = time.time()
             self.assertRaises(PinException, self._some_protected_call, False, True, False)
             test_backoff(attempt, start)
+
 '''
         # Unplug Trezor now
         self.client.debuglink.stop()