diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 06b1eba..00d562e 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2491,39 +2491,39 @@
 
 PSA key derivation: invalid type (0)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid type (PSA_KEY_TYPE_CATEGORY_MASK)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid length (0)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:0:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid length (7 bits)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:7:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:7:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: raw data, 8 bits
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:8:PSA_SUCCESS
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:8:PSA_SUCCESS:0
 
 PSA key derivation: invalid length (9 bits)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:9:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:9:PSA_ERROR_INVALID_ARGUMENT:0
 
 # This test assumes that PSA_MAX_KEY_BITS (currently 65536-8 bits = 8191 bytes
 # and not expected to be raised any time soon) is less than the maximum
 # output from HKDF-SHA512 (255*64 = 16320 bytes).
 PSA key derivation: largest possible key
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512
-derive_large_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS:PSA_SUCCESS
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS:PSA_SUCCESS:1
 
 PSA key derivation: key too large
 depends_on:MBEDTLS_SHA512_C
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: good
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
@@ -2649,113 +2649,113 @@
 generate_random:2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1
 
 PSA generate key: bad type (0)
-generate_key:PSA_KEY_TYPE_NONE:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_NONE:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: bad type (PSA_KEY_TYPE_CATEGORY_MASK)
-generate_key:PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: bad type (RSA public key)
-generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: raw data, 0 bits: invalid argument
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-generate_key:PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, 7 bits: invalid argument
-generate_key:PSA_KEY_TYPE_RAW_DATA:7:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:7:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
 PSA generate key: raw data, 9 bits: invalid argument
-generate_key:PSA_KEY_TYPE_RAW_DATA:9:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:9:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, (MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:(MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:(MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
 PSA generate key: raw data, (2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:(2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:(2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
-PSA generate key: raw data, 65528 bits (ok)
-generate_large_key:PSA_KEY_TYPE_RAW_DATA:65528:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+PSA generate key: raw data, 65528 bits (large key, ok if it fits)
+generate_key:PSA_KEY_TYPE_RAW_DATA:65528:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:1
 
 PSA generate key: raw data, 65536 bits (not supported)
-generate_key:PSA_KEY_TYPE_RAW_DATA:65536:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RAW_DATA:65536:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: AES, 128 bits, CTR
 depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR
-generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_SUCCESS:0
 
 PSA generate key: AES, 128 bits, GCM
 depends_on:MBEDTLS_AES_C:MBEDTLS_GCM_C
-generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_GCM:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_GCM:PSA_SUCCESS:0
 
 PSA generate key: DES, 64 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: DES, 128 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: DES, 192 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: invalid key size: AES, 64 bits
 depends_on:MBEDTLS_AES_C
-generate_key:PSA_KEY_TYPE_AES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_AES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: RSA, 512 bits, good, sign (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1016 bits, good, sign (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1016:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1016:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1024 bits, good, sign (PSS SHA-256)
 depends_on:PSA_WANT_ALG_RSA_PSS:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):PSA_SUCCESS:0
 
 PSA generate key: RSA, 512 bits, good, encrypt (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1024 bits, good, encrypt (OAEP SHA-256)
 depends_on:PSA_WANT_ALG_RSA_OAEP:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256):PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256):PSA_SUCCESS:0
 
 PSA generate key: RSA, 0 bits: invalid
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: RSA, 1022 bits: not supported
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1022:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1022:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: RSA, 1023 bits: not supported
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1023:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1023:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: RSA, maximum size exceeded
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_MAX_KEY_BITS+1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_MAX_KEY_BITS+1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: ECC, SECP256R1, good
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_SECP256R1_ENABLED
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
 
 PSA generate key: ECC, SECP256R1, incorrect bit size
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDSA_C
 # INVALID_ARGUMENT would make more sense, but our code as currently structured
 # doesn't fully relate the curve with its size.
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: ECC, Curve25519, good
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_CURVE25519_ENABLED:!MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):255:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):255:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS:0
 
 PSA generate key: RSA, default e
 generate_key_rsa:512:"":PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 061c52e..1051e0c 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1592,6 +1592,7 @@
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
     psa_set_key_type( &attributes, type );
     status = psa_import_key( &attributes, buffer, byte_size, &key );
+    TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
     TEST_EQUAL( status, expected_status );
 
     if( status == PSA_SUCCESS )
@@ -5284,7 +5285,8 @@
 void derive_key( int alg_arg,
                  data_t *key_data, data_t *input1, data_t *input2,
                  int type_arg, int bits_arg,
-                 int expected_status_arg )
+                 int expected_status_arg,
+                 int is_large_output )
 {
     mbedtls_svc_key_id_t base_key = MBEDTLS_SVC_KEY_ID_INIT;
     mbedtls_svc_key_id_t derived_key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -5313,9 +5315,14 @@
     psa_set_key_algorithm( &derived_attributes, 0 );
     psa_set_key_type( &derived_attributes, type );
     psa_set_key_bits( &derived_attributes, bits );
-    TEST_EQUAL( psa_key_derivation_output_key( &derived_attributes, &operation,
-                                               &derived_key ),
-                expected_status );
+
+    psa_status_t status =
+      psa_key_derivation_output_key( &derived_attributes,
+                                     &operation,
+                                     &derived_key );
+    if( is_large_output > 0 )
+      TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
+    TEST_EQUAL( status, expected_status );
 
 exit:
     psa_key_derivation_abort( &operation );
@@ -5326,26 +5333,6 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void derive_large_key( int alg_arg,
-                       data_t *key_data, data_t *input1, data_t *input2,
-                       int type_arg, int bits_arg,
-                       int expected_status_arg )
-{
-    size_t key_bytes = PSA_BITS_TO_BYTES(bits_arg);
-    uint8_t* buffer = NULL;
-
-    /* Check that the target running this test can accomodate large
-     * keys on its heap, before calling the actual generate_key test */
-    ASSERT_ALLOC_WEAK(buffer, key_bytes);
-    mbedtls_free( buffer );
-
-    test_derive_key( alg_arg, key_data, input1, input2, type_arg, bits_arg, expected_status_arg );
-exit:
-    return;
-}
-/* END_CASE */
-
-/* BEGIN_CASE */
 void key_agreement_setup( int alg_arg,
                           int our_key_type_arg, int our_key_alg_arg,
                           data_t *our_key_data, data_t *peer_key_data,
@@ -5609,7 +5596,8 @@
                    int bits_arg,
                    int usage_arg,
                    int alg_arg,
-                   int expected_status_arg )
+                   int expected_status_arg,
+                   int is_large_key )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t type = type_arg;
@@ -5628,7 +5616,11 @@
     psa_set_key_bits( &attributes, bits );
 
     /* Generate a key */
-    TEST_EQUAL( psa_generate_key( &attributes, &key ), expected_status );
+    psa_status_t status = psa_generate_key( &attributes, &key );
+
+    if( is_large_key > 0 )
+      TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
+    TEST_EQUAL( status , expected_status );
     if( expected_status != PSA_SUCCESS )
         goto exit;
 
@@ -5653,27 +5645,6 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE */
-void generate_large_key( int type_arg,
-                         int bits_arg,
-                         int usage_arg,
-                         int alg_arg,
-                         int expected_status_arg )
-{
-    size_t key_bytes = PSA_BITS_TO_BYTES(bits_arg);
-    uint8_t* buffer = NULL;
-
-    /* Check that the target running this test can accomodate large
-     * keys on its heap, before calling the actual generate_key test */
-    ASSERT_ALLOC_WEAK(buffer, key_bytes);
-    mbedtls_free( buffer );
-
-    test_generate_key( type_arg, bits_arg, usage_arg, alg_arg, expected_status_arg );
-exit:
-    return;
-}
-/* END_CASE */
-
 /* BEGIN_CASE depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME:MBEDTLS_PKCS1_V15 */
 void generate_key_rsa( int bits_arg,
                        data_t *e_arg,