From 9e1f590d5a929a21223bb39ae5123d27e69fda66 Mon Sep 17 00:00:00 2001
From: Szczepan Zalega <szczepan@nitrokey.com>
Date: Fri, 17 Jan 2020 18:19:45 +0100
Subject: Add tests for Pro`s bootloader

---
 unittest/constants.py |  3 +++
 unittest/test_pro.py  | 31 +++++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/unittest/constants.py b/unittest/constants.py
index 645ef6a..bf4f5b9 100644
--- a/unittest/constants.py
+++ b/unittest/constants.py
@@ -39,6 +39,9 @@ class DefaultPasswords:
     USER_TEMP = b'234234234'
     UPDATE = b'12345678'
     UPDATE_TEMP = b'123update123'
+    UPDATE_LONG = b'1234567890'*2
+    UPDATE_TOO_LONG = UPDATE_LONG + b'x'
+    UPDATE_TOO_SHORT = b'1234567'
 
 
 class DeviceErrorCode:
diff --git a/unittest/test_pro.py b/unittest/test_pro.py
index a8df7cd..3f1cb8d 100644
--- a/unittest/test_pro.py
+++ b/unittest/test_pro.py
@@ -977,21 +977,48 @@ def test_get_device_model(C):
     # assert C.NK_get_device_model() != C.NK_DISCONNECTED
 
 
+@pytest.mark.firmware
+def test_bootloader_password_change_pro_length(C):
+    skip_if_device_version_lower_than({'P': 11})
+
+    # Test whether the correct password is set
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Change to the longest possible password
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_LONG) == DeviceErrorCode.STATUS_OK
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_LONG, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Use longer or shorter passwords than possible
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TOO_LONG) == LibraryErrors.TOO_LONG_STRING
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TOO_SHORT) == DeviceErrorCode.WRONG_PASSWORD
+
+
+
 @pytest.mark.firmware
 def test_bootloader_password_change_pro(C):
     skip_if_device_version_lower_than({'P': 11})
     assert C.NK_change_firmware_password_pro(b'zxcasd', b'zxcasd') == DeviceErrorCode.WRONG_PASSWORD
 
+    # Revert effects of broken test run, if needed
+    C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_TEMP, DefaultPasswords.UPDATE)
+
+    # Change to the same password
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    # Change password
     assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE, DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.STATUS_OK
     assert C.NK_change_firmware_password_pro(DefaultPasswords.UPDATE_TEMP, DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
 
 
 @pytest.mark.firmware
-def test_bootloader_run_pro(C):
+def test_bootloader_run_pro_wrong_password(C):
     skip_if_device_version_lower_than({'P': 11})
     assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE_TEMP) == DeviceErrorCode.WRONG_PASSWORD
+
+
+@pytest.mark.skip
+@pytest.mark.firmware
+def test_bootloader_run_pro(C):
     # Not enabled due to lack of side-effect removal at this point
-    # assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
+    assert C.NK_enable_firmware_update_pro(DefaultPasswords.UPDATE) == DeviceErrorCode.STATUS_OK
 
 
 @pytest.mark.firmware
-- 
cgit v1.2.1