diff --git a/tf-psa-crypto/include/psa/crypto_extra.h b/tf-psa-crypto/include/psa/crypto_extra.h
index f48c087..9bc68f6 100644
--- a/tf-psa-crypto/include/psa/crypto_extra.h
+++ b/tf-psa-crypto/include/psa/crypto_extra.h
@@ -42,6 +42,13 @@
      PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE : PSA_CIPHER_MAX_KEY_LENGTH)
 #endif /* !MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE*/
 
+/* Define the size of the each key slot buffer. */
+#if defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+#define MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE     MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE
+#else
+#define MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE     SIZE_MAX
+#endif
+
 /** \addtogroup attributes
  * @{
  */
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
index fffe498..f67961f 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
@@ -7158,7 +7158,7 @@
 # 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:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= 16320
 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
@@ -7402,15 +7402,15 @@
 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
-depends_on:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= (MBEDTLS_CTR_DRBG_MAX_REQUEST + 1)
 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
-depends_on:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= (2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1)
 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 (large key, ok if it fits)
-depends_on:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= 65528
 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)
@@ -7633,15 +7633,15 @@
 concurrently_generate_keys:PSA_KEY_TYPE_RAW_DATA:9:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0:8:5
 
 PSA concurrent key generation: raw data, (MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-depends_on:MBEDTLS_THREADING_PTHREAD:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_THREADING_PTHREAD:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= (MBEDTLS_CTR_DRBG_MAX_REQUEST + 1)
 concurrently_generate_keys:PSA_KEY_TYPE_RAW_DATA:(MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0:8:5
 
 PSA concurrent key generation: raw data, (2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-depends_on:MBEDTLS_THREADING_PTHREAD:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_THREADING_PTHREAD:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE >= (2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1)
 concurrently_generate_keys:PSA_KEY_TYPE_RAW_DATA:(2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0:8:5
 
 PSA concurrent key generation: raw data, 65528 bits (large key, ok if it fits)
-depends_on:MBEDTLS_THREADING_PTHREAD:!MBEDTLS_PSA_STATIC_KEY_SLOTS
+depends_on:MBEDTLS_THREADING_PTHREAD:MBEDTLS_PSA_KEY_BUFFER_MAX_SIZE > 65528
 concurrently_generate_keys:PSA_KEY_TYPE_RAW_DATA:65528:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:1:8:5
 
 PSA concurrent key generation: raw data, 65536 bits (not supported)