Fix mypy on ragger tests

LedgerHQ · Feb 26, 2024 · 2d69e1b · 2d69e1b
1 parent ec50977
commit 2d69e1b
Showing 1 changed file with 8 additions and 11 deletions.
diff --git a/tests/test_seed.py b/tests/test_seed.py
@@ -6,7 +6,6 @@
 """
 import sys
 import pytest
-from Crypto.PublicKey import RSA, ECC
 
 from application_client.command_sender import CommandSender
 from application_client.app_def import Errors, DataObject, PassWord, PubkeyAlgo
@@ -15,7 +14,7 @@
 from utils import check_pincode, generate_key, get_key_attributes, KEY_TEMPLATES
 
 
-def _gen_key(client: CommandSender, template: str) -> RSA.RsaKey | ECC.EccKey:
+def _gen_key(client: CommandSender, template: str):
 
     # Verify PW3 (Admin)
     check_pincode(client, PassWord.PW3)
@@ -31,15 +30,13 @@ def _gen_key(client: CommandSender, template: str) -> RSA.RsaKey | ECC.EccKey:
 
      # Read the SIG pub Key
     if key_algo == PubkeyAlgo.RSA:
-        pubkey: RSA.RsaKey = get_RSA_pub_key(client, DataObject.DO_SIG_KEY)
-    elif key_algo == PubkeyAlgo.ECDSA:
-        pubkey: ECC.EccKey = get_ECDSA_pub_key(client, DataObject.DO_SIG_KEY)
-    elif key_algo == PubkeyAlgo.EDDSA:
-        pubkey: ECC.EccKey = get_EDDSA_pub_key(client, DataObject.DO_SIG_KEY)
-    else:
-        raise ValueError
-
-    return pubkey
+        return get_RSA_pub_key(client, DataObject.DO_SIG_KEY)
+    if key_algo == PubkeyAlgo.ECDSA:
+        return get_ECDSA_pub_key(client, DataObject.DO_SIG_KEY)
+    if key_algo == PubkeyAlgo.EDDSA:
+        return get_EDDSA_pub_key(client, DataObject.DO_SIG_KEY)
+
+    raise ValueError
 
 
 @pytest.mark.parametrize(