diff --git a/tests/scripts/generate_test_keys.py b/tests/scripts/generate_test_keys.py
index 1236baa..0a67a78 100755
--- a/tests/scripts/generate_test_keys.py
+++ b/tests/scripts/generate_test_keys.py
@@ -80,10 +80,12 @@
         return ""
     return prefix + str(bits) + suffix
 
-def get_look_up_table_entry(key_type: str, curve_or_keybits: str,
+def get_look_up_table_entry(key_type: str, group_id_or_keybits: str,
                             priv_array_name: str, pub_array_name: str) -> Iterator[str]:
-    yield "    {{ {}, ".format("1" if key_type == "ec" else "0")
-    yield "{},\n".format(curve_or_keybits)
+    if key_type == "ec":
+        yield "    {{ {}, 0,\n".format(group_id_or_keybits)
+    else:
+        yield "    {{ 0, {},\n".format(group_id_or_keybits)
     yield "      {0}, sizeof({0}),\n".format(priv_array_name)
     yield "      {0}, sizeof({0}) }},".format(pub_array_name)
 
@@ -135,16 +137,16 @@
             output_file.write(''.join(["\n", c_array_priv, "\n", c_array_pub, "\n"]))
             # Update the lookup table
             if key_type == "ec":
-                curve_or_keybits = "MBEDTLS_ECP_DP_" + curve.upper()
+                group_id_or_keybits = "MBEDTLS_ECP_DP_" + curve.upper()
             else:
-                curve_or_keybits = str(bits)
-            look_up_table.append(''.join(get_look_up_table_entry(key_type, curve_or_keybits,
+                group_id_or_keybits = str(bits)
+            look_up_table.append(''.join(get_look_up_table_entry(key_type, group_id_or_keybits,
                                                                  array_name_priv, array_name_pub)))
     # Write the lookup table: the struct containing pointers to all the arrays we created above.
     output_file.write("""
 struct predefined_key_element {
-    int is_ec;  // 1 for EC keys; 0 for RSA
-    int curve_or_keybits;
+    int group_id;  // EC group ID; 0 for RSA keys
+    int keybits;  // bits size of RSA key; 0 for EC keys
     const unsigned char *priv_key;
     size_t priv_key_len;
     const unsigned char *pub_key;
diff --git a/tests/src/test_keys.h b/tests/src/test_keys.h
index 7e8f773..ec54fe4 100644
--- a/tests/src/test_keys.h
+++ b/tests/src/test_keys.h
@@ -730,8 +730,8 @@
 };
 
 struct predefined_key_element {
-    int is_ec;  // 1 for EC keys; 0 for RSA
-    int curve_or_keybits;
+    int group_id;  // EC group ID; 0 for RSA keys
+    int keybits;  // bits size of RSA key; 0 for EC keys
     const unsigned char *priv_key;
     size_t priv_key_len;
     const unsigned char *pub_key;
@@ -739,40 +739,40 @@
 };
 
 struct predefined_key_element predefined_keys[] = {
-    { 1, MBEDTLS_ECP_DP_SECP192K1,
+    { MBEDTLS_ECP_DP_SECP192K1, 0,
       test_ec_secp192k1_priv, sizeof(test_ec_secp192k1_priv),
       test_ec_secp192k1_pub, sizeof(test_ec_secp192k1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP256K1,
+    { MBEDTLS_ECP_DP_SECP256K1, 0,
       test_ec_secp256k1_priv, sizeof(test_ec_secp256k1_priv),
       test_ec_secp256k1_pub, sizeof(test_ec_secp256k1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP192R1,
+    { MBEDTLS_ECP_DP_SECP192R1, 0,
       test_ec_secp192r1_priv, sizeof(test_ec_secp192r1_priv),
       test_ec_secp192r1_pub, sizeof(test_ec_secp192r1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP224R1,
+    { MBEDTLS_ECP_DP_SECP224R1, 0,
       test_ec_secp224r1_priv, sizeof(test_ec_secp224r1_priv),
       test_ec_secp224r1_pub, sizeof(test_ec_secp224r1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP256R1,
+    { MBEDTLS_ECP_DP_SECP256R1, 0,
       test_ec_secp256r1_priv, sizeof(test_ec_secp256r1_priv),
       test_ec_secp256r1_pub, sizeof(test_ec_secp256r1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP384R1,
+    { MBEDTLS_ECP_DP_SECP384R1, 0,
       test_ec_secp384r1_priv, sizeof(test_ec_secp384r1_priv),
       test_ec_secp384r1_pub, sizeof(test_ec_secp384r1_pub) },
-    { 1, MBEDTLS_ECP_DP_SECP521R1,
+    { MBEDTLS_ECP_DP_SECP521R1, 0,
       test_ec_secp521r1_priv, sizeof(test_ec_secp521r1_priv),
       test_ec_secp521r1_pub, sizeof(test_ec_secp521r1_pub) },
-    { 1, MBEDTLS_ECP_DP_BP256R1,
+    { MBEDTLS_ECP_DP_BP256R1, 0,
       test_ec_bp256r1_priv, sizeof(test_ec_bp256r1_priv),
       test_ec_bp256r1_pub, sizeof(test_ec_bp256r1_pub) },
-    { 1, MBEDTLS_ECP_DP_BP384R1,
+    { MBEDTLS_ECP_DP_BP384R1, 0,
       test_ec_bp384r1_priv, sizeof(test_ec_bp384r1_priv),
       test_ec_bp384r1_pub, sizeof(test_ec_bp384r1_pub) },
-    { 1, MBEDTLS_ECP_DP_BP512R1,
+    { MBEDTLS_ECP_DP_BP512R1, 0,
       test_ec_bp512r1_priv, sizeof(test_ec_bp512r1_priv),
       test_ec_bp512r1_pub, sizeof(test_ec_bp512r1_pub) },
-    { 1, MBEDTLS_ECP_DP_CURVE25519,
+    { MBEDTLS_ECP_DP_CURVE25519, 0,
       test_ec_curve25519_priv, sizeof(test_ec_curve25519_priv),
       test_ec_curve25519_pub, sizeof(test_ec_curve25519_pub) },
-    { 1, MBEDTLS_ECP_DP_CURVE448,
+    { MBEDTLS_ECP_DP_CURVE448, 0,
       test_ec_curve448_priv, sizeof(test_ec_curve448_priv),
       test_ec_curve448_pub, sizeof(test_ec_curve448_pub) },
     { 0, 1024,
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index bee3273..ad7da32 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -186,24 +186,33 @@
 #define RSA_KEY_SIZE   1024
 #define RSA_KEY_LEN   (RSA_KEY_SIZE/8)
 
-static int get_predefined_key_data(int is_ec, int curve_or_keybits,
+static int get_predefined_key_data(int is_ec, int group_id_or_keybits,
                                    const unsigned char **key, size_t *key_len,
                                    const unsigned char **pub_key, size_t *pub_key_len)
 {
     size_t i;
+    struct predefined_key_element *predefined_key = NULL;
+
     for (i = 0; i < ARRAY_LENGTH(predefined_keys); i++) {
-        if ((is_ec == predefined_keys[i].is_ec) &&
-            (curve_or_keybits == predefined_keys[i].curve_or_keybits)) {
-            *key = predefined_keys[i].priv_key;
-            *key_len = predefined_keys[i].priv_key_len;
-            if (pub_key != NULL) {
-                *pub_key = predefined_keys[i].pub_key;
-                *pub_key_len = predefined_keys[i].pub_key_len;
+        if (is_ec) {
+            if (group_id_or_keybits == predefined_keys[i].group_id) {
+                predefined_key = &predefined_keys[i];
             }
-            return 0;
+        } else if (group_id_or_keybits == predefined_keys[i].keybits) {
+            predefined_key = &predefined_keys[i];
         }
     }
 
+    if (predefined_key != NULL) {
+        *key = predefined_key->priv_key;
+        *key_len = predefined_key->priv_key_len;
+        if (pub_key != NULL) {
+            *pub_key = predefined_key->pub_key;
+            *pub_key_len = predefined_key->pub_key_len;
+        }
+        return 0;
+    }
+
     TEST_FAIL("Unsupported key");
     /* "exit" label is to make the compiler happy. */
 exit: