diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index e697491..af1e20c 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -621,62 +621,62 @@
 depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME:MBEDTLS_PK_WRITE_C
 pk_psa_sign:1024:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA256
+PK sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA256
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA384
+PK sign ext: RSA2048, PK_RSA, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA384
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA512
+PK sign ext: RSA2048, PK_RSA, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA512
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
 
-PK Sign ext:SECP256R1,PK_ECDSA,MD_SHA256
+PK sign ext: SECP256R1, PK_ECDSA, MD_SHA256
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256R1:MBEDTLS_MD_CAN_SHA256
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA256
 
-PK Sign ext:SECP384R1,PK_ECDSA,MD_SHA384
+PK sign ext: SECP384R1, PK_ECDSA, MD_SHA384
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP384R1:MBEDTLS_MD_CAN_SHA384
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA384
 
-PK Sign ext:SECP521R1,PK_ECDSA,MD_SHA512
+PK sign ext: SECP521R1, PK_ECDSA, MD_SHA512
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP521R1:MBEDTLS_MD_CAN_SHA512
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA512
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA256
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA256
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA384
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA384
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA512
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA512
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index fa0b03b..5255feb 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -407,6 +407,16 @@
                                            buf, buf_size, &buf_size,
                                            NULL, NULL,
                                            NULL));
+    TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
+               mbedtls_pk_sign_ext(pk_type, &ctx, MBEDTLS_MD_NONE,
+                                   NULL, buf_size,
+                                   buf, buf_size, &buf_size,
+                                   NULL, NULL));
+    TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
+               mbedtls_pk_sign_ext(pk_type, &ctx, MBEDTLS_MD_SHA256,
+                                   NULL, 0,
+                                   buf, buf_size, &buf_size,
+                                   NULL, NULL));
 exit:
     mbedtls_pk_free(&ctx);
     USE_PSA_DONE();
@@ -443,14 +453,6 @@
                                             NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
-    TEST_ASSERT(mbedtls_pk_sign_restartable(&pk,
-                                            MBEDTLS_MD_NONE,
-                                            NULL, 0,
-                                            buf, sizeof(buf), &len,
-                                            mbedtls_test_rnd_std_rand, NULL,
-                                            NULL) ==
-                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
-
     TEST_ASSERT(mbedtls_pk_sign(&pk,
                                 MBEDTLS_MD_NONE,
                                 NULL, 0,
@@ -458,6 +460,13 @@
                                 mbedtls_test_rnd_std_rand, NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
+    TEST_ASSERT(mbedtls_pk_sign_ext(MBEDTLS_PK_NONE, &pk,
+                                    MBEDTLS_MD_NONE,
+                                    NULL, 0,
+                                    buf, sizeof(buf), &len,
+                                    mbedtls_test_rnd_std_rand, NULL) ==
+                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+
     TEST_ASSERT(mbedtls_pk_verify_restartable(&pk,
                                               MBEDTLS_MD_NONE,
                                               NULL, 0,
@@ -1175,22 +1184,31 @@
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                 mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk,
+                                mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)), 0);
 
 #if defined(MBEDTLS_PKCS1_V21)
-    TEST_ASSERT(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, NULL, &pk,
-                                      MBEDTLS_MD_NONE, hash, hash_len, sig, sig_len) ==
-                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+    TEST_EQUAL(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, NULL, &pk,
+                                     MBEDTLS_MD_NONE, hash, hash_len, sig, sig_len),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 #endif /* MBEDTLS_PKCS1_V21 */
 
-    TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_NONE, hash, hash_len,
-                                  sig, sig_len) == MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+    TEST_EQUAL(mbedtls_pk_verify(&pk, MBEDTLS_MD_NONE, hash, hash_len,
+                                 sig, sig_len),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
-    TEST_ASSERT(mbedtls_pk_sign(&pk, MBEDTLS_MD_NONE, hash, hash_len,
-                                sig, sizeof(sig), &sig_len,
-                                mbedtls_test_rnd_std_rand, NULL)
-                == MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+#if defined(MBEDTLS_PKCS1_V21)
+    TEST_EQUAL(mbedtls_pk_sign_ext(MBEDTLS_PK_RSASSA_PSS, &pk,
+                                   MBEDTLS_MD_NONE, hash, hash_len,
+                                   sig, sizeof(sig), &sig_len,
+                                   mbedtls_test_rnd_std_rand, NULL),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+#endif /* MBEDTLS_PKCS1_V21 */
+
+    TEST_EQUAL(mbedtls_pk_sign(&pk, MBEDTLS_MD_NONE, hash, hash_len,
+                               sig, sizeof(sig), &sig_len,
+                               mbedtls_test_rnd_std_rand, NULL),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
 exit:
     mbedtls_pk_free(&pk);
@@ -1440,14 +1458,14 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_GENPRIME */
-void pk_psa_sign_ext(int pk_type, int parameter, int key_pk_type, int md_alg)
+/* BEGIN_CASE depends_on:MBEDTLS_GENPRIME */
+void pk_sign_ext(int pk_type, int parameter, int key_pk_type, int md_alg)
 {
     /* See the description of pk_genkey() for the description of the `parameter` argument. */
     mbedtls_pk_context pk;
     size_t sig_len;
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    unsigned char hash[PSA_HASH_MAX_SIZE];
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
     size_t hash_len = mbedtls_md_get_size_from_type(md_alg);
     void const *options = NULL;
     mbedtls_pk_rsassa_pss_options rsassa_pss_options;
@@ -1455,16 +1473,15 @@
     memset(sig, 0, sizeof(sig));
 
     mbedtls_pk_init(&pk);
-    PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                 mbedtls_pk_info_from_type(pk_type)) == 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk,
+                                mbedtls_pk_info_from_type(pk_type)), 0);
+    TEST_EQUAL(pk_genkey(&pk, parameter), 0);
 
-    TEST_ASSERT(pk_genkey(&pk, parameter) == 0);
-
-    TEST_ASSERT(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
-                                    sig, sizeof(sig), &sig_len,
-                                    mbedtls_test_rnd_std_rand, NULL) == 0);
+    TEST_EQUAL(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
+                                   sig, sizeof(sig), &sig_len,
+                                   mbedtls_test_rnd_std_rand, NULL), 0);
 
     if (key_pk_type == MBEDTLS_PK_RSASSA_PSS) {
         rsassa_pss_options.mgf1_hash_id = md_alg;
@@ -1472,11 +1489,11 @@
         rsassa_pss_options.expected_salt_len = hash_len;
         options = (const void *) &rsassa_pss_options;
     }
-    TEST_ASSERT(mbedtls_pk_verify_ext(key_pk_type, options, &pk, md_alg,
-                                      hash, hash_len, sig, sig_len) == 0);
+    TEST_EQUAL(mbedtls_pk_verify_ext(key_pk_type, options, &pk, md_alg,
+                                     hash, hash_len, sig, sig_len), 0);
 exit:
     mbedtls_pk_free(&pk);
-    PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */