diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 5ee491e..f214394 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -476,6 +476,13 @@
 
 #endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
 
+/* Helper macro to get 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
+
 /* Helper macro for the PK module to check whether MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE
  * is large enough to contain 4096 bits RSA key pairs. Of course this check is only
  * necessary if PK relies on PSA (i.e. MBEDTLS_USE_PSA_CRYPTO) to store and manage
diff --git a/tf-psa-crypto/include/psa/crypto_extra.h b/tf-psa-crypto/include/psa/crypto_extra.h
index 9bc68f6..f48c087 100644
--- a/tf-psa-crypto/include/psa/crypto_extra.h
+++ b/tf-psa-crypto/include/psa/crypto_extra.h
@@ -42,13 +42,6 @@
      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
  * @{
  */