diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 9168b1d..47d72d0 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -74,17 +74,17 @@
  * or PK_PSA_INVALID_SLOT if no slot was available.
  * The key uses NIST P-256 and is usable for signing with SHA-256.
  */
-psa_key_slot_t pk_psa_genkey( void )
+psa_key_handle_t pk_psa_genkey( void )
 {
-    psa_key_slot_t key;
+    psa_key_handle_t key;
 
     const int curve = PSA_ECC_CURVE_SECP256R1;
     const psa_key_type_t type = PSA_KEY_TYPE_ECC_KEYPAIR(curve);
     const size_t bits = 256;
     psa_key_policy_t policy;
 
-    /* find a free key slot */
-    if( PSA_SUCCESS != mbedtls_psa_get_free_key_slot( &key ) )
+    /* Allocate a key slot */
+    if( PSA_SUCCESS != psa_allocate_key( type, bits, &key ) )
         return( PK_PSA_INVALID_SLOT );
 
     /* set up policy on key slot */
@@ -112,7 +112,7 @@
 void pk_psa_utils(  )
 {
     mbedtls_pk_context pk, pk2;
-    psa_key_slot_t key;
+    psa_key_handle_t key;
 
     const char * const name = "Opaque";
     const size_t bitlen = 256; /* harcoded in genkey() */
@@ -778,7 +778,7 @@
 void pk_psa_sign(  )
 {
     mbedtls_pk_context pk;
-    psa_key_slot_t key;
+    psa_key_handle_t key;
     unsigned char hash[50], sig[100], pkey[100];
     size_t sig_len, klen = 0;
 
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index bf43a80..268b4bb 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -133,7 +133,7 @@
                                  int cert_type )
 {
     mbedtls_pk_context key;
-    psa_key_slot_t slot;
+    psa_key_handle_t slot;
     psa_algorithm_t md_alg_psa;
     mbedtls_x509write_csr req;
     unsigned char buf[4096];