diff --git a/tests/include/test/drivers/crypto_config_test_driver_extension.h b/tests/include/test/drivers/crypto_config_test_driver_extension.h
index dac07acd33..66378e7def 100644
--- a/tests/include/test/drivers/crypto_config_test_driver_extension.h
+++ b/tests/include/test/drivers/crypto_config_test_driver_extension.h
@@ -1,9 +1,24 @@
 /**
- * This file is intended to be used to build PSA test driver libraries. It is
- * intended to be appended by the test build system to the crypto_config.h file
- * of the Mbed TLS library the test library will be linked to. It mirrors the
- * PSA_ACCEL_* macros defining the cryptographic operations the test library
- * supports.
+ * This file is intended to be used to build PSA external test driver
+ * libraries (libtestdriver1).
+ *
+ * It is intended to be appended by the test build system to the
+ * crypto_config.h file of the Mbed TLS library the test library will be
+ * linked to (see `tests/Makefile` libtestdriver1 target). This is done in
+ * order to insert it at the right time: after the main configuration
+ * (PSA_WANT) but before the logic that determines what built-ins to enable
+ * based on PSA_WANT and MBEDTLS_PSA_ACCEL macros.
+ *
+ * It reverses the PSA_ACCEL_* macros defining the cryptographic operations
+ * that will be accelerated in the main library:
+ * - When something is accelerated in the main library, we need it supported
+ *   in libtestdriver1, so we disable the accel macro in order to the built-in
+ *   to be enabled.
+ * - When something is NOT accelerated in the main library, we don't need it
+ *   in libtestdriver1, so we enable its accel macro in order to the built-in
+ *   to be disabled, to keep libtestdriver1 minimal. (We can't adjust the
+ *   PSA_WANT macros as they need to be the same between libtestdriver1 and
+ *   the main library, since they determine the ABI between the two.)
  */
 
 #include "psa/crypto_legacy.h"
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 1039712aef..233dbe6ad3 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -11,7 +11,7 @@
 
 #include "test/helpers.h"
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
 #include "test/psa_helpers.h"
 #include <psa/crypto.h>
 #endif
@@ -40,12 +40,15 @@
         mbedtls_psa_crypto_free();                                      \
     }                                                                   \
     while (0)
-#else /*MBEDTLS_PSA_CRYPTO_C */
+#elif defined(MBEDTLS_PSA_CRYPTO_CLIENT) /* MBEDTLS_PSA_CRYPTO_CLIENT && !MBEDTLS_PSA_CRYPTO_C */
+#define PSA_INIT() PSA_ASSERT(psa_crypto_init())
+#define PSA_DONE() mbedtls_psa_crypto_free();
+#else  /* MBEDTLS_PSA_CRYPTO_CLIENT && !MBEDTLS_PSA_CRYPTO_C */
 #define PSA_INIT() ((void) 0)
 #define PSA_DONE() ((void) 0)
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
 
@@ -293,18 +296,19 @@ uint64_t mbedtls_test_parse_binary_string(data_t *bin_string);
  *  \param  alg             The AEAD algorithm.
  *  \param  nonce_length    The nonce length in number of bytes.
  */
+
 #if defined(MBEDTLS_GCM_ALT) || \
     defined(MBEDTLS_PSA_ACCEL_ALG_GCM)
-#define MBEDTLS_TEST_HAVE_ALT_GCM  1
+#define MBEDTLS_TEST_HAVE_ACCEL_GCM  1
 #else
-#define MBEDTLS_TEST_HAVE_ALT_GCM  0
+#define MBEDTLS_TEST_HAVE_ACCEL_GCM  0
 #endif
 
 #define MBEDTLS_TEST_PSA_SKIP_IF_ALT_GCM_NOT_12BYTES_NONCE(alg,           \
                                                            nonce_length) \
     do                                                                     \
     {                                                                      \
-        if ((MBEDTLS_TEST_HAVE_ALT_GCM) &&                               \
+        if ((MBEDTLS_TEST_HAVE_ACCEL_GCM) &&                               \
             (PSA_ALG_AEAD_WITH_SHORTENED_TAG((alg), 0) ==            \
              PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM, 0)) &&       \
             ((nonce_length) != 12))                                   \
@@ -315,7 +319,7 @@ uint64_t mbedtls_test_parse_binary_string(data_t *bin_string);
     }                                                                      \
     while (0)
 
-#endif /* MBEDTLS_PSA_CRYPTO_C */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
 
 /** \def USE_PSA_INIT
  *
@@ -334,18 +338,9 @@ uint64_t mbedtls_test_parse_binary_string(data_t *bin_string);
  * This is like #PSA_DONE except it does nothing under the same conditions as
  * #USE_PSA_INIT.
  */
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_SSL_PROTO_TLS1_3)
 #define USE_PSA_INIT() PSA_INIT()
 #define USE_PSA_DONE() PSA_DONE()
-#elif defined(MBEDTLS_SSL_PROTO_TLS1_3)
-/* TLS 1.3 must work without having called psa_crypto_init(), for backward
- * compatibility with Mbed TLS <= 3.5 when connecting with a peer that
- * supports both TLS 1.2 and TLS 1.3. See mbedtls_ssl_tls13_crypto_init()
- * and https://github.com/Mbed-TLS/mbedtls/issues/9072 . */
-#define USE_PSA_INIT() ((void) 0)
-/* TLS 1.3 may have initialized the PSA subsystem. Shut it down cleanly,
- * otherwise Asan and Valgrind would notice a resource leak. */
-#define USE_PSA_DONE() PSA_DONE()
 #else /* MBEDTLS_USE_PSA_CRYPTO || MBEDTLS_SSL_PROTO_TLS1_3 */
 /* Define empty macros so that we can use them in the preamble and teardown
  * of every test function that uses PSA conditionally based on
@@ -417,12 +412,13 @@ uint64_t mbedtls_test_parse_binary_string(data_t *bin_string);
  * This is like #PSA_DONE except it does nothing under the same conditions as
  * #MD_OR_USE_PSA_INIT.
  */
-#if defined(MBEDTLS_MD_SOME_PSA)
+#if defined(MBEDTLS_MD_SOME_PSA) || \
+    defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_SSL_PROTO_TLS1_3)
 #define MD_OR_USE_PSA_INIT()   PSA_INIT()
 #define MD_OR_USE_PSA_DONE()   PSA_DONE()
 #else
-#define MD_OR_USE_PSA_INIT()   USE_PSA_INIT()
-#define MD_OR_USE_PSA_DONE()   USE_PSA_DONE()
+#define MD_OR_USE_PSA_INIT() ((void) 0)
+#define MD_OR_USE_PSA_DONE() ((void) 0)
 #endif
 
 /** \def AES_PSA_INIT
diff --git a/tests/include/test/psa_test_wrappers.h b/tests/include/test/psa_test_wrappers.h
index 134a547c85..7ab2bea6b9 100644
--- a/tests/include/test/psa_test_wrappers.h
+++ b/tests/include/test/psa_test_wrappers.h
@@ -363,13 +363,22 @@ psa_status_t mbedtls_test_wrap_psa_generate_key_custom(
 #define psa_generate_key_custom(arg0_attributes, arg1_custom, arg2_custom_data, arg3_custom_data_length, arg4_key) \
     mbedtls_test_wrap_psa_generate_key_custom(arg0_attributes, arg1_custom, arg2_custom_data, arg3_custom_data_length, arg4_key)
 
-psa_status_t mbedtls_test_wrap_psa_generate_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    const psa_key_production_parameters_t *arg1_params,
-    size_t arg2_params_data_length,
-    mbedtls_svc_key_id_t *arg3_key);
-#define psa_generate_key_ext(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key) \
-    mbedtls_test_wrap_psa_generate_key_ext(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key)
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_abort(
+    psa_generate_key_iop_t *arg0_operation);
+#define psa_generate_key_iop_abort(arg0_operation) \
+    mbedtls_test_wrap_psa_generate_key_iop_abort(arg0_operation)
+
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
+    psa_generate_key_iop_t *arg0_operation,
+    psa_key_id_t *arg1_key);
+#define psa_generate_key_iop_complete(arg0_operation, arg1_key) \
+    mbedtls_test_wrap_psa_generate_key_iop_complete(arg0_operation, arg1_key)
+
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_setup(
+    psa_generate_key_iop_t *arg0_operation,
+    const psa_key_attributes_t *arg1_attributes);
+#define psa_generate_key_iop_setup(arg0_operation, arg1_attributes) \
+    mbedtls_test_wrap_psa_generate_key_iop_setup(arg0_operation, arg1_attributes)
 
 psa_status_t mbedtls_test_wrap_psa_generate_random(
     uint8_t *arg0_output,
@@ -449,6 +458,37 @@ psa_status_t mbedtls_test_wrap_psa_import_key(
 #define psa_import_key(arg0_attributes, arg1_data, arg2_data_length, arg3_key) \
     mbedtls_test_wrap_psa_import_key(arg0_attributes, arg1_data, arg2_data_length, arg3_key)
 
+psa_status_t mbedtls_test_wrap_psa_key_agreement(
+    mbedtls_svc_key_id_t arg0_private_key,
+    const uint8_t *arg1_peer_key,
+    size_t arg2_peer_key_length,
+    psa_algorithm_t arg3_alg,
+    const psa_key_attributes_t *arg4_attributes,
+    mbedtls_svc_key_id_t *arg5_key);
+#define psa_key_agreement(arg0_private_key, arg1_peer_key, arg2_peer_key_length, arg3_alg, arg4_attributes, arg5_key) \
+    mbedtls_test_wrap_psa_key_agreement(arg0_private_key, arg1_peer_key, arg2_peer_key_length, arg3_alg, arg4_attributes, arg5_key)
+
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_abort(
+    psa_key_agreement_iop_t *arg0_operation);
+#define psa_key_agreement_iop_abort(arg0_operation) \
+    mbedtls_test_wrap_psa_key_agreement_iop_abort(arg0_operation)
+
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_complete(
+    psa_key_agreement_iop_t *arg0_operation,
+    psa_key_id_t *arg1_key);
+#define psa_key_agreement_iop_complete(arg0_operation, arg1_key) \
+    mbedtls_test_wrap_psa_key_agreement_iop_complete(arg0_operation, arg1_key)
+
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_setup(
+    psa_key_agreement_iop_t *arg0_operation,
+    psa_key_id_t arg1_private_key,
+    const uint8_t *arg2_peer_key,
+    size_t arg3_peer_key_length,
+    psa_algorithm_t arg4_alg,
+    const psa_key_attributes_t *arg5_attributes);
+#define psa_key_agreement_iop_setup(arg0_operation, arg1_private_key, arg2_peer_key, arg3_peer_key_length, arg4_alg, arg5_attributes) \
+    mbedtls_test_wrap_psa_key_agreement_iop_setup(arg0_operation, arg1_private_key, arg2_peer_key, arg3_peer_key_length, arg4_alg, arg5_attributes)
+
 psa_status_t mbedtls_test_wrap_psa_key_derivation_abort(
     psa_key_derivation_operation_t *arg0_operation);
 #define psa_key_derivation_abort(arg0_operation) \
@@ -515,15 +555,6 @@ psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_custom(
 #define psa_key_derivation_output_key_custom(arg0_attributes, arg1_operation, arg2_custom, arg3_custom_data, arg4_custom_data_length, arg5_key) \
     mbedtls_test_wrap_psa_key_derivation_output_key_custom(arg0_attributes, arg1_operation, arg2_custom, arg3_custom_data, arg4_custom_data_length, arg5_key)
 
-psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    psa_key_derivation_operation_t *arg1_operation,
-    const psa_key_production_parameters_t *arg2_params,
-    size_t arg3_params_data_length,
-    mbedtls_svc_key_id_t *arg4_key);
-#define psa_key_derivation_output_key_ext(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key) \
-    mbedtls_test_wrap_psa_key_derivation_output_key_ext(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key)
-
 psa_status_t mbedtls_test_wrap_psa_key_derivation_set_capacity(
     psa_key_derivation_operation_t *arg0_operation,
     size_t arg1_capacity);
diff --git a/tests/src/drivers/hash.c b/tests/src/drivers/hash.c
index 76ec12a22f..54aec93225 100644
--- a/tests/src/drivers/hash.c
+++ b/tests/src/drivers/hash.c
@@ -13,7 +13,11 @@
 #include "test/drivers/hash.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_hash.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_hash.h"
+#endif
 #endif
 
 mbedtls_test_driver_hash_hooks_t
diff --git a/tests/src/drivers/test_driver_aead.c b/tests/src/drivers/test_driver_aead.c
index 314ce83a25..6992a066d2 100644
--- a/tests/src/drivers/test_driver_aead.c
+++ b/tests/src/drivers/test_driver_aead.c
@@ -16,7 +16,11 @@
 #include "mbedtls/constant_time.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_aead.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_aead.h"
+#endif
 #endif
 
 mbedtls_test_driver_aead_hooks_t
diff --git a/tests/src/drivers/test_driver_asymmetric_encryption.c b/tests/src/drivers/test_driver_asymmetric_encryption.c
index 4fc8c9d34b..6fdbe43aee 100644
--- a/tests/src/drivers/test_driver_asymmetric_encryption.c
+++ b/tests/src/drivers/test_driver_asymmetric_encryption.c
@@ -16,7 +16,11 @@
 #include "test/drivers/key_management.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_rsa.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_rsa.h"
+#endif
 #endif
 
 #define PSA_RSA_KEY_PAIR_MAX_SIZE \
diff --git a/tests/src/drivers/test_driver_cipher.c b/tests/src/drivers/test_driver_cipher.c
index 2bc751a8a2..90256fc4ea 100644
--- a/tests/src/drivers/test_driver_cipher.c
+++ b/tests/src/drivers/test_driver_cipher.c
@@ -19,7 +19,11 @@
 #include "test/random.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_cipher.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_cipher.h"
+#endif
 #endif
 
 #include <string.h>
diff --git a/tests/src/drivers/test_driver_key_agreement.c b/tests/src/drivers/test_driver_key_agreement.c
index 8471959e2a..8a7a9ea525 100644
--- a/tests/src/drivers/test_driver_key_agreement.c
+++ b/tests/src/drivers/test_driver_key_agreement.c
@@ -20,9 +20,15 @@
 #include <string.h>
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/include/psa/crypto.h"
 #include "libtestdriver1/library/psa_crypto_ecp.h"
 #include "libtestdriver1/library/psa_crypto_ffdh.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/include/psa/crypto.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_ffdh.h"
+#endif
 #endif
 
 mbedtls_test_driver_key_agreement_hooks_t
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index 2a878994c2..d2ca157b8b 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -23,9 +23,15 @@
 #include "test/random.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_ecp.h"
 #include "libtestdriver1/library/psa_crypto_rsa.h"
 #include "libtestdriver1/library/psa_crypto_ffdh.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_rsa.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_ffdh.h"
+#endif
 #endif
 
 #include <string.h>
diff --git a/tests/src/drivers/test_driver_mac.c b/tests/src/drivers/test_driver_mac.c
index 9f8120bd4a..f1cf504303 100644
--- a/tests/src/drivers/test_driver_mac.c
+++ b/tests/src/drivers/test_driver_mac.c
@@ -13,7 +13,11 @@
 #include "test/drivers/mac.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_mac.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_mac.h"
+#endif
 #endif
 
 mbedtls_test_driver_mac_hooks_t mbedtls_test_driver_mac_hooks =
diff --git a/tests/src/drivers/test_driver_pake.c b/tests/src/drivers/test_driver_pake.c
index a0b6c1cb0c..c3ce326fe2 100644
--- a/tests/src/drivers/test_driver_pake.c
+++ b/tests/src/drivers/test_driver_pake.c
@@ -1,5 +1,5 @@
 /*
- * Test driver for MAC entry points.
+ * Test driver for PAKE entry points.
  */
 /*  Copyright The Mbed TLS Contributors
  *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
@@ -14,7 +14,11 @@
 #include "string.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_pake.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_pake.h"
+#endif
 #endif
 
 mbedtls_test_driver_pake_hooks_t mbedtls_test_driver_pake_hooks =
diff --git a/tests/src/drivers/test_driver_signature.c b/tests/src/drivers/test_driver_signature.c
index 4fca5d178d..a6eef573b2 100644
--- a/tests/src/drivers/test_driver_signature.c
+++ b/tests/src/drivers/test_driver_signature.c
@@ -26,9 +26,15 @@
 #include "test/random.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#if MBEDTLS_VERSION_MAJOR < 4
 #include "libtestdriver1/library/psa_crypto_ecp.h"
 #include "libtestdriver1/library/psa_crypto_hash.h"
 #include "libtestdriver1/library/psa_crypto_rsa.h"
+#else
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_hash.h"
+#include "libtestdriver1/tf-psa-crypto/drivers/builtin/src/psa_crypto_rsa.h"
+#endif
 #endif
 
 #include <string.h>
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 937bd45d22..a496de0212 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -11,7 +11,9 @@
 #include <test/macros.h>
 #include <test/psa_exercise_key.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
+#if ((MBEDTLS_VERSION_MAJOR < 4) \
+    && defined(MBEDTLS_PSA_CRYPTO_C)) \
+    || defined(MBEDTLS_PSA_CRYPTO_CLIENT)
 
 #include <mbedtls/asn1.h>
 #include <psa/crypto.h>
@@ -1284,7 +1286,7 @@ int mbedtls_test_key_consistency_psa_pk(mbedtls_svc_key_id_t psa_key,
             break;
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
@@ -1295,7 +1297,7 @@ int mbedtls_test_key_consistency_psa_pk(mbedtls_svc_key_id_t psa_key,
                            pk_public_buffer, sizeof(pk_public_buffer)), 0);
             pk_public = pk_public_buffer;
             break;
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS && !MBEDTLS_PK_USE_PSA_EC_DATA */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY && !MBEDTLS_PK_USE_PSA_EC_DATA */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
         case MBEDTLS_PK_OPAQUE:
@@ -1332,4 +1334,4 @@ exit:
 }
 #endif /* MBEDTLS_PK_C */
 
-#endif /* MBEDTLS_PSA_CRYPTO_C */
+#endif /* MBEDTLS_PSA_CRYPTO_C || MBEDTLS_PSA_CRYPTO_CLIENT */
diff --git a/tests/src/psa_test_wrappers.c b/tests/src/psa_test_wrappers.c
index eceb40bc70..6e42a8bba6 100644
--- a/tests/src/psa_test_wrappers.c
+++ b/tests/src/psa_test_wrappers.c
@@ -622,14 +622,29 @@ psa_status_t mbedtls_test_wrap_psa_generate_key_custom(
     return status;
 }
 
-/* Wrapper for psa_generate_key_ext */
-psa_status_t mbedtls_test_wrap_psa_generate_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    const psa_key_production_parameters_t *arg1_params,
-    size_t arg2_params_data_length,
-    mbedtls_svc_key_id_t *arg3_key)
+/* Wrapper for psa_generate_key_iop_abort */
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_abort(
+    psa_generate_key_iop_t *arg0_operation)
 {
-    psa_status_t status = (psa_generate_key_ext)(arg0_attributes, arg1_params, arg2_params_data_length, arg3_key);
+    psa_status_t status = (psa_generate_key_iop_abort)(arg0_operation);
+    return status;
+}
+
+/* Wrapper for psa_generate_key_iop_complete */
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
+    psa_generate_key_iop_t *arg0_operation,
+    psa_key_id_t *arg1_key)
+{
+    psa_status_t status = (psa_generate_key_iop_complete)(arg0_operation, arg1_key);
+    return status;
+}
+
+/* Wrapper for psa_generate_key_iop_setup */
+psa_status_t mbedtls_test_wrap_psa_generate_key_iop_setup(
+    psa_generate_key_iop_t *arg0_operation,
+    const psa_key_attributes_t *arg1_attributes)
+{
+    psa_status_t status = (psa_generate_key_iop_setup)(arg0_operation, arg1_attributes);
     return status;
 }
 
@@ -790,6 +805,61 @@ psa_status_t mbedtls_test_wrap_psa_import_key(
     return status;
 }
 
+/* Wrapper for psa_key_agreement */
+psa_status_t mbedtls_test_wrap_psa_key_agreement(
+    mbedtls_svc_key_id_t arg0_private_key,
+    const uint8_t *arg1_peer_key,
+    size_t arg2_peer_key_length,
+    psa_algorithm_t arg3_alg,
+    const psa_key_attributes_t *arg4_attributes,
+    mbedtls_svc_key_id_t *arg5_key)
+{
+#if !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS)
+    MBEDTLS_TEST_MEMORY_POISON(arg1_peer_key, arg2_peer_key_length);
+#endif /* !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS) */
+    psa_status_t status = (psa_key_agreement)(arg0_private_key, arg1_peer_key, arg2_peer_key_length, arg3_alg, arg4_attributes, arg5_key);
+#if !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS)
+    MBEDTLS_TEST_MEMORY_UNPOISON(arg1_peer_key, arg2_peer_key_length);
+#endif /* !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS) */
+    return status;
+}
+
+/* Wrapper for psa_key_agreement_iop_abort */
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_abort(
+    psa_key_agreement_iop_t *arg0_operation)
+{
+    psa_status_t status = (psa_key_agreement_iop_abort)(arg0_operation);
+    return status;
+}
+
+/* Wrapper for psa_key_agreement_iop_complete */
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_complete(
+    psa_key_agreement_iop_t *arg0_operation,
+    psa_key_id_t *arg1_key)
+{
+    psa_status_t status = (psa_key_agreement_iop_complete)(arg0_operation, arg1_key);
+    return status;
+}
+
+/* Wrapper for psa_key_agreement_iop_setup */
+psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_setup(
+    psa_key_agreement_iop_t *arg0_operation,
+    psa_key_id_t arg1_private_key,
+    const uint8_t *arg2_peer_key,
+    size_t arg3_peer_key_length,
+    psa_algorithm_t arg4_alg,
+    const psa_key_attributes_t *arg5_attributes)
+{
+#if !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS)
+    MBEDTLS_TEST_MEMORY_POISON(arg2_peer_key, arg3_peer_key_length);
+#endif /* !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS) */
+    psa_status_t status = (psa_key_agreement_iop_setup)(arg0_operation, arg1_private_key, arg2_peer_key, arg3_peer_key_length, arg4_alg, arg5_attributes);
+#if !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS)
+    MBEDTLS_TEST_MEMORY_UNPOISON(arg2_peer_key, arg3_peer_key_length);
+#endif /* !defined(MBEDTLS_PSA_ASSUME_EXCLUSIVE_BUFFERS) */
+    return status;
+}
+
 /* Wrapper for psa_key_derivation_abort */
 psa_status_t mbedtls_test_wrap_psa_key_derivation_abort(
     psa_key_derivation_operation_t *arg0_operation)
@@ -907,18 +977,6 @@ psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_custom(
     return status;
 }
 
-/* Wrapper for psa_key_derivation_output_key_ext */
-psa_status_t mbedtls_test_wrap_psa_key_derivation_output_key_ext(
-    const psa_key_attributes_t *arg0_attributes,
-    psa_key_derivation_operation_t *arg1_operation,
-    const psa_key_production_parameters_t *arg2_params,
-    size_t arg3_params_data_length,
-    mbedtls_svc_key_id_t *arg4_key)
-{
-    psa_status_t status = (psa_key_derivation_output_key_ext)(arg0_attributes, arg1_operation, arg2_params, arg3_params_data_length, arg4_key);
-    return status;
-}
-
 /* Wrapper for psa_key_derivation_set_capacity */
 psa_status_t mbedtls_test_wrap_psa_key_derivation_set_capacity(
     psa_key_derivation_operation_t *arg0_operation,