diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 5cc924c..cac86ca 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -2382,19 +2382,6 @@
     mbedtls_pk_init(&pk_pub);
     PSA_INIT();
 
-    /* Get the MD type to be used for the tests below from the provided key policy. */
-    mbedtls_md_type_t md_for_test = MBEDTLS_MD_SHA256; /* Default */
-    int is_psa_hash_alg_specified = 0;
-    if ((PSA_ALG_GET_HASH(key_alg) != PSA_ALG_NONE) &&
-        (PSA_ALG_GET_HASH(key_alg) != PSA_ALG_ANY_HASH)) {
-        md_for_test = mbedtls_md_type_from_psa_alg(key_alg);
-        is_psa_hash_alg_specified = 1;
-    }
-
-    in_buf_len = mbedtls_md_get_size_from_type(md_for_test);
-    TEST_CALLOC(in_buf, in_buf_len);
-    memset(in_buf, 0x1, in_buf_len);
-
     if (key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
         key_usage |= PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT;
     }
@@ -2439,6 +2426,22 @@
     /* Test #2: check that the 2 generated PK contexts form a valid private/public key pair. */
     TEST_EQUAL(mbedtls_pk_check_pair(&pk_pub, &pk_priv, mbedtls_test_rnd_std_rand, NULL), 0);
 
+    /* Get the MD alg to be used for the tests below from the provided key policy. */
+    mbedtls_md_type_t md_for_test = MBEDTLS_MD_SHA256; /* Default */
+    if ((PSA_ALG_GET_HASH(key_alg) != PSA_ALG_NONE) &&
+        (PSA_ALG_GET_HASH(key_alg) != PSA_ALG_ANY_HASH)) {
+        md_for_test = mbedtls_md_type_from_psa_alg(key_alg);
+    }
+    /* Use also the same MD algorithm for PSA sign/verify checks. This is helpful
+     * for the cases in which the key policy algorithm is ANY_HASH type. */
+    psa_algorithm_t psa_alg_for_test =
+        (key_alg & ~PSA_ALG_HASH_MASK) |
+        (mbedtls_md_psa_alg_from_type(md_for_test) & PSA_ALG_HASH_MASK);
+
+    in_buf_len = mbedtls_md_get_size_from_type(md_for_test);
+    TEST_CALLOC(in_buf, in_buf_len);
+    memset(in_buf, 0x1, in_buf_len);
+
     /* Test #3: sign/verify with the following pattern:
      * - Sign using the PK context generated from the private key.
      * - Verify from the same PK context used for signature.
@@ -2454,7 +2457,7 @@
     TEST_EQUAL(mbedtls_pk_verify(&pk_pub, md_for_test, in_buf, in_buf_len,
                                  out_buf, out_buf_len), 0);
 
-    if (PSA_ALG_IS_SIGN(key_alg) && is_psa_hash_alg_specified) {
+    if (PSA_ALG_IS_HASH_AND_SIGN(key_alg)) {
 #if defined(MBEDTLS_PSA_UTIL_HAVE_ECDSA)
         /* ECDSA signature requires PK->PSA format conversion. */
         if (PSA_ALG_IS_ECDSA(key_alg)) {
@@ -2463,7 +2466,7 @@
                                                 sizeof(out_buf), &out_buf_len), 0);
         }
 #endif /* MBEDTLS_PSA_UTIL_HAVE_ECDSA */
-        PSA_ASSERT(psa_verify_hash(pub_key_id, key_alg, in_buf, in_buf_len,
+        PSA_ASSERT(psa_verify_hash(pub_key_id, psa_alg_for_test, in_buf, in_buf_len,
                                    out_buf, out_buf_len));
     }
 
@@ -2471,8 +2474,8 @@
      * sign with PSA and verify with PK. Key's policy must include a valid hash
      * algorithm (not any).
      */
-    if (PSA_ALG_IS_SIGN(key_alg) && is_psa_hash_alg_specified) {
-        PSA_ASSERT(psa_sign_hash(priv_key_id, key_alg, in_buf, in_buf_len,
+    if (PSA_ALG_IS_HASH_AND_SIGN(key_alg)) {
+        PSA_ASSERT(psa_sign_hash(priv_key_id, psa_alg_for_test, in_buf, in_buf_len,
                                  out_buf, sizeof(out_buf), &out_buf_len));
 #if defined(MBEDTLS_PSA_UTIL_HAVE_ECDSA)
         /*  ECDSA signature requires PSA->PK format conversion */