From 949438f4c5eee9fee82252e20e8a91790c8bb004 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Thu, 6 Sep 2018 19:05:26 +0200
Subject: [PATCH] tests: add script_type test to getpublickey test

---
 .../device_tests/test_msg_getpublickey.py     | 36 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/trezorlib/tests/device_tests/test_msg_getpublickey.py b/trezorlib/tests/device_tests/test_msg_getpublickey.py
index 09b5dfcf66..287d496327 100644
--- a/trezorlib/tests/device_tests/test_msg_getpublickey.py
+++ b/trezorlib/tests/device_tests/test_msg_getpublickey.py
@@ -14,7 +14,7 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
-from trezorlib import btc
+from trezorlib import btc, messages as proto
 from trezorlib.tools import H_
 
 from ..support import ckd_public as bip32
@@ -132,3 +132,37 @@ class TestMsgGetpublickey(TrezorTest):
             btc.get_public_node(self.client, [111, 42], coin_name="Testnet").xpub
             == "tpubDAgixSyai5PWbc8N1mBkHDR5nLgAnHFtY7r4y5EzxqAxrt9YUDpZL3kaRoHVvCfrcwNo31c2isBP2uTHcZxEosuKbyJhCAbrvGoPuLUZ7Mz"
         )
+
+    def test_script_type(self):
+        self.setup_mnemonic_nopin_nopassphrase()
+        assert (
+            btc.get_public_node(self.client, [], coin_name="Bitcoin").xpub
+            == "xpub661MyMwAqRbcF1zGijBb2K6x9YiJPh58xpcCeLvTxMX6spkY3PcpJ4ABcCyWfskq5DDxM3e6Ez5ePCqG5bnPUXR4wL8TZWyoDaUdiWW7bKy"
+        )
+        assert (
+            btc.get_public_node(
+                self.client,
+                [],
+                coin_name="Bitcoin",
+                script_type=proto.InputScriptType.SPENDADDRESS,
+            ).xpub
+            == "xpub661MyMwAqRbcF1zGijBb2K6x9YiJPh58xpcCeLvTxMX6spkY3PcpJ4ABcCyWfskq5DDxM3e6Ez5ePCqG5bnPUXR4wL8TZWyoDaUdiWW7bKy"
+        )
+        assert (
+            btc.get_public_node(
+                self.client,
+                [],
+                coin_name="Bitcoin",
+                script_type=proto.InputScriptType.SPENDP2SHWITNESS,
+            ).xpub
+            == "ypub6QqdH2c5z7966KBPZ5yDEQCTKWrkLK4dsw8RRjpMLMtyvvZmJ3nNv7pKdQw6fnQkUrLm6XEeheSCGVSpoJCQGm6fofpt9RoHVJYH72ecmVm"
+        )
+        assert (
+            btc.get_public_node(
+                self.client,
+                [],
+                coin_name="Bitcoin",
+                script_type=proto.InputScriptType.SPENDWITNESS,
+            ).xpub
+            == "zpub6jftahH18ngZwcNWPSkqSVHxVV1CGw48o3eeD8iEiNGrz2NzYhwwYBUTectgfh4ftVTZqzqDAJnk9n4PWzcR4znGg1XJjLcmm2bvVc3Honv"
+        )