diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index de90b47..fd923c2 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -69,37 +69,26 @@
 
 #include "mbedtls/psa_util.h"
 
-#define PK_PSA_INVALID_SLOT 0 /* guaranteed invalid */
-
 /*
- * Generate a key in a free key slot and return this key slot,
- * or PK_PSA_INVALID_SLOT if no slot was available.
+ * Generate a key using PSA and return a handle to that key,
+ * or 0 if the key generation failed.
  * The key uses NIST P-256 and is usable for signing with SHA-256.
  */
 psa_key_handle_t pk_psa_genkey( void )
 {
     psa_key_handle_t key;
-
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     const int curve = PSA_ECC_CURVE_SECP256R1;
     const psa_key_type_t type = PSA_KEY_TYPE_ECC_KEY_PAIR(curve);
     const size_t bits = 256;
-    psa_key_policy_t policy;
 
-    /* Allocate a key slot */
-    if( PSA_SUCCESS != psa_allocate_key( &key ) )
-        return( PK_PSA_INVALID_SLOT );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN );
+    psa_set_key_algorithm( &attributes, PSA_ALG_ECDSA(PSA_ALG_SHA_256) );
+    psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
+    PSA_ASSERT( psa_generate_key( &attributes, &key ) );
 
-    /* set up policy on key slot */
-    policy = psa_key_policy_init();
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN,
-                                      PSA_ALG_ECDSA(PSA_ALG_SHA_256) );
-    if( PSA_SUCCESS != psa_set_key_policy( key, &policy ) )
-        return( PK_PSA_INVALID_SLOT );
-
-    /* generate key */
-    if( PSA_SUCCESS != psa_generate_key_to_handle( key, type, bits, NULL, 0 ) )
-        return( PK_PSA_INVALID_SLOT );
-
+exit:
     return( key );
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
@@ -115,6 +104,7 @@
 {
     mbedtls_pk_context pk, pk2;
     psa_key_handle_t key;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     const char * const name = "Opaque";
     const size_t bitlen = 256; /* harcoded in genkey() */
@@ -136,7 +126,8 @@
     mbedtls_pk_init( &pk );
 
     key = pk_psa_genkey();
-    TEST_ASSERT( key != 0 );
+    if( key == 0 )
+        goto exit;
 
     TEST_ASSERT( mbedtls_pk_setup_opaque( &pk, key ) == 0 );
 
@@ -173,7 +164,7 @@
 
     /* test that freeing the context does not destroy the key */
     mbedtls_pk_free( &pk );
-    TEST_ASSERT( PSA_SUCCESS == psa_get_key_information( key, NULL, NULL ) );
+    TEST_ASSERT( PSA_SUCCESS == psa_get_key_attributes( key, &attributes ) );
     TEST_ASSERT( PSA_SUCCESS == psa_destroy_key( key ) );
 
 exit:
@@ -1233,7 +1224,6 @@
     pkey_legacy_start = pkey_legacy + sizeof( pkey_legacy ) - klen_legacy;
 
     /* Turn PK context into an opaque one. */
-    TEST_ASSERT( psa_allocate_key( &handle ) == PSA_SUCCESS );
     TEST_ASSERT( mbedtls_pk_wrap_as_opaque( &pk, &handle,
                                             PSA_ALG_SHA_256 ) == 0 );