Adapt to the new key allocation mechanism

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;