diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 3779aa9..4806d09 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -227,28 +227,26 @@
  * (mbedtls_pk_setup() has been called on the PK context ) so that it
  * can determine the key type to be loaded from the PK context itself.
  *
- * \param pk               The PK object to fill. It must have been initialized
- *                         with mbedtls_pk_setup().
+ * \param pk               The PK object to fill. It must  have been initialized
+ *                         (mbedtls_pk_init()), but not setup (mbedtls_pk_setup()).
+ * \param pk_info          mbedtls_pk_info_t to use in the generated PK context.
  * \param curve_or_keybits - For RSA keys, the key size in bits.
  *                         - For EC keys, the curve (\c MBEDTLS_ECP_DP_xxx).
  *
  * \return                 0 on success or a negative value otherwise.
  */
-static int pk_genkey(mbedtls_pk_context *pk, int curve_or_keybits)
+static int pk_genkey(mbedtls_pk_context *pk, const mbedtls_pk_info_t *pk_info, int curve_or_keybits)
 {
     unsigned char *key_data = NULL;
     size_t key_data_len = 0;
     int ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    int is_rsa = (mbedtls_pk_get_type(pk) == MBEDTLS_PK_RSA);
-
-    /* Dump pk_info since this is overridden by mbedtls_pk_parse_keyfile(). */
-    const mbedtls_pk_info_t *original_pk_info = pk->pk_info;
+    int is_rsa = (curve_or_keybits >= 1024);
 
     TEST_EQUAL(get_predefined_key_data(is_rsa, curve_or_keybits, &key_data, &key_data_len), 0);
     TEST_EQUAL(mbedtls_pk_parse_key(pk, key_data, key_data_len, NULL, 0,
                                     mbedtls_test_rnd_std_rand, NULL), 0);
-    /* Restore pk_info. */
-    pk->pk_info = original_pk_info;
+    /* Override pk_info. */
+    pk->pk_info = pk_info;
     ret = 0;
 
 exit:
@@ -311,7 +309,7 @@
      * Note: mbedtls_pk_write_key_der() and mbedtls_mpi_write_binary() write
      * key data at the end of the provided buffer, whereas psa_export_key()
      * writes the key at the beginning.
-    */
+     */
     if (mbedtls_pk_get_type(&pk) == MBEDTLS_PK_RSA) {
 #if defined(MBEDTLS_PK_WRITE_C)
         ret = mbedtls_pk_write_key_der(&pk, key_data, key_data_size);
@@ -499,12 +497,12 @@
     if (pk_type == MBEDTLS_PK_NONE) {
         return 0;
     }
-    TEST_EQUAL(mbedtls_pk_setup(pk, mbedtls_pk_info_from_type(pk_type)), 0);
 
     switch (pk_type) {
 #if defined(MBEDTLS_RSA_C)
         case MBEDTLS_PK_RSA:
         {
+            TEST_EQUAL(mbedtls_pk_setup(pk, mbedtls_pk_info_from_type(pk_type)), 0);
             *psa_type = PSA_KEY_TYPE_RSA_KEY_PAIR;
             mbedtls_rsa_context *rsa = mbedtls_pk_rsa(*pk);
             if (want_pair) {
@@ -538,7 +536,7 @@
             mbedtls_ecp_group_id grp_id = MBEDTLS_TEST_ECP_DP_ONE_CURVE;
             size_t bits;
             *psa_type = PSA_KEY_TYPE_ECC_KEY_PAIR(mbedtls_ecc_group_to_psa(grp_id, &bits));
-            TEST_EQUAL(pk_genkey(pk, grp_id), 0);
+            TEST_EQUAL(pk_genkey(pk, mbedtls_pk_info_from_type(pk_type), grp_id), 0);
             if (!want_pair) {
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
                 psa_key_attributes_t pub_attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -799,9 +797,7 @@
 
         TEST_EQUAL(mbedtls_pk_get_type(&pk), MBEDTLS_PK_OPAQUE);
     } else {
-        TEST_EQUAL(mbedtls_pk_setup(&pk,
-                                    mbedtls_pk_info_from_type(key_type)), 0);
-        TEST_EQUAL(pk_genkey(&pk, curve_or_keybits), 0);
+        TEST_EQUAL(pk_genkey(&pk, mbedtls_pk_info_from_type(key_type), curve_or_keybits), 0);
         TEST_EQUAL(mbedtls_pk_get_type(&pk), key_type);
     }
 
@@ -1003,8 +999,7 @@
     mbedtls_pk_init(&pk);
     USE_PSA_INIT();
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
-    TEST_ASSERT(pk_genkey(&pk, curve_or_keybits) == 0);
+    TEST_ASSERT(pk_genkey(&pk, mbedtls_pk_info_from_type(type), curve_or_keybits) == 0);
 
     TEST_ASSERT((int) mbedtls_pk_get_type(&pk) == type);
     TEST_ASSERT(mbedtls_pk_can_do(&pk, type));
@@ -1372,8 +1367,7 @@
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
-    TEST_ASSERT(pk_genkey(&pk, curve_or_keybits) == 0);
+    TEST_ASSERT(pk_genkey(&pk, mbedtls_pk_info_from_type(type), curve_or_keybits) == 0);
 
 #if defined(MBEDTLS_RSA_C)
     if (type == MBEDTLS_PK_RSA) {
@@ -1766,9 +1760,7 @@
     memset(test, 0, sizeof(test));
 
     /* Initialize PK RSA context with random key */
-    TEST_ASSERT(mbedtls_pk_setup(&rsa,
-                                 mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
-    TEST_ASSERT(pk_genkey(&rsa, RSA_KEY_SIZE) == 0);
+    TEST_ASSERT(pk_genkey(&rsa, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA), RSA_KEY_SIZE) == 0);
 
     /* Extract key to the raw rsa context */
     TEST_ASSERT(mbedtls_rsa_copy(&raw, mbedtls_pk_rsa(rsa)) == 0);
@@ -1869,9 +1861,7 @@
     /* Create the legacy EC/RSA PK context. */
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
-        TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                     mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
-        TEST_EQUAL(pk_genkey(&pk, bits), 0);
+        TEST_EQUAL(pk_genkey(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA), bits), 0);
         TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);
     }
 #else /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
@@ -1880,8 +1870,7 @@
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         ecp_grp_id = mbedtls_ecc_group_from_psa(psa_type, bits);
-        TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
-        TEST_ASSERT(pk_genkey(&pk, ecp_grp_id) == 0);
+        TEST_ASSERT(pk_genkey(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY), ecp_grp_id) == 0);
     }
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
@@ -2019,9 +2008,7 @@
     mbedtls_pk_init(&pk);
     MD_OR_USE_PSA_INIT();
 
-    TEST_EQUAL(mbedtls_pk_setup(&pk,
-                                mbedtls_pk_info_from_type(pk_type)), 0);
-    TEST_EQUAL(pk_genkey(&pk, curve_or_keybits), 0);
+    TEST_EQUAL(pk_genkey(&pk, mbedtls_pk_info_from_type(pk_type), curve_or_keybits), 0);
 
     TEST_EQUAL(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
                                    sig, sizeof(sig), &sig_len,