Update persistent_key_load_key_from_storage to use attributes

Update persistent_key_load_key_from_storage to the new attribute-based
key creation interface. I tweaked the code a little to make it simpler
and more robust without changing the core logic.
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index e148dd9..3392f64 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2036,34 +2036,42 @@
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDSA_C
 generate_key:PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:PSA_ERROR_INVALID_ARGUMENT
 
-persistent key can be accessed after in-memory deletion: AES, 128 bits, CTR
+PSA import persistent key: raw data, 0 bits
 depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT:PSA_ALG_CTR:IMPORT_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:IMPORT_KEY
 
-PSA generate persistent key: raw data, 8 bits
+PSA import persistent key: AES, 128 bits, exportable
+depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
+persistent_key_load_key_from_storage:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_EXPORT:PSA_ALG_CTR:IMPORT_KEY
+
+PSA import persistent key: AES, 128 bits, non-exportable
+depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
+persistent_key_load_key_from_storage:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CTR:IMPORT_KEY
+
+PSA generate persistent key: raw data, 8 bits, exportable
 depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:GENERATE_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:GENERATE_KEY
 
-PSA generate persistent key: AES, 128 bits, CTR
+PSA generate persistent key: AES, 128 bits, exportable
 depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:GENERATE_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:GENERATE_KEY
 
-PSA generate persistent key: DES, 64 bits, CBC-nopad
+PSA generate persistent key: AES, 128 bits, non-exportable
+depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:MBEDTLS_PSA_CRYPTO_STORAGE_C
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:GENERATE_KEY
+
+PSA generate persistent key: DES, 64 bits, exportable
 depends_on:MBEDTLS_DES_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:GENERATE_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:GENERATE_KEY
 
-PSA generate persistent key: RSA, 1024 bits, good, sign (PSS SHA-256)
+PSA generate persistent key: RSA, 1024 bits, exportable
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME:MBEDTLS_PKCS1_V21:MBEDTLS_SHA256_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_RSA_KEYPAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):GENERATE_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_RSA_KEYPAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):GENERATE_KEY
 
-PSA generate persistent key: ECC, SECP256R1, good
+PSA generate persistent key: ECC, SECP256R1, exportable
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDSA_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:GENERATE_KEY:PSA_SUCCESS
+persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:GENERATE_KEY
 
-PSA derive persistent key: HKDF SHA-256
+PSA derive persistent key: HKDF SHA-256, exportable
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_TYPE_RAW_DATA:1024:PSA_KEY_USAGE_EXPORT:0:DERIVE_KEY:PSA_SUCCESS
-
-PSA generate persistent key: AES, 128 bits, CTR
-depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:MBEDTLS_PSA_CRYPTO_STORAGE_C
-persistent_key_load_key_from_storage:"":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:GENERATE_KEY:PSA_ERROR_NOT_PERMITTED
+persistent_key_load_key_from_storage:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_TYPE_RAW_DATA:1024:PSA_KEY_USAGE_EXPORT:0:DERIVE_KEY
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b1964a4..e656c64 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1098,7 +1098,7 @@
     TEST_EQUAL( psa_get_key_attributes( handle, &attributes ),
                 PSA_ERROR_INVALID_HANDLE );
     TEST_EQUAL( psa_get_key_id( &attributes ), 0 );
-    TEST_EQUAL( psa_get_key_attributes_lifetime( &attributes ), 0 );
+    TEST_EQUAL( psa_get_key_lifetime( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_type( &attributes ), 0 );
@@ -4715,22 +4715,19 @@
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
-void persistent_key_load_key_from_storage( data_t *data, int type_arg,
-                                           int bits, int usage_arg,
-                                           int alg_arg, int generation_method,
-                                           int export_status )
+void persistent_key_load_key_from_storage( data_t *data,
+                                           int type_arg, int bits_arg,
+                                           int usage_flags_arg, int alg_arg,
+                                           int generation_method )
 {
+    psa_key_id_t key_id = 1;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_handle_t handle = 0;
-    psa_key_handle_t base_key;
-    psa_key_type_t type = (psa_key_type_t) type_arg;
-    psa_key_type_t type_get;
-    size_t bits_get;
-    psa_key_policy_t policy_set = PSA_KEY_POLICY_INIT;
-    psa_key_policy_t policy_get = PSA_KEY_POLICY_INIT;
-    psa_key_usage_t policy_usage = (psa_key_usage_t) usage_arg;
-    psa_algorithm_t policy_alg = (psa_algorithm_t) alg_arg;
-    psa_key_policy_t base_policy_set = PSA_KEY_POLICY_INIT;
-    psa_algorithm_t base_policy_alg = PSA_ALG_HKDF(PSA_ALG_SHA_256);
+    psa_key_handle_t base_key = 0;
+    psa_key_type_t type = type_arg;
+    size_t bits = bits_arg;
+    psa_key_usage_t usage_flags = usage_flags_arg;
+    psa_algorithm_t alg = alg_arg;
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
     unsigned char *first_export = NULL;
     unsigned char *second_export = NULL;
@@ -4738,102 +4735,115 @@
     size_t first_exported_length;
     size_t second_exported_length;
 
-    ASSERT_ALLOC( first_export, export_size );
-    ASSERT_ALLOC( second_export, export_size );
+    if( usage_flags & PSA_KEY_USAGE_EXPORT )
+    {
+        ASSERT_ALLOC( first_export, export_size );
+        ASSERT_ALLOC( second_export, export_size );
+    }
 
     PSA_ASSERT( psa_crypto_init() );
 
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, 1,
-                                &handle ) );
-    psa_key_policy_set_usage( &policy_set, policy_usage,
-                              policy_alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy_set ) );
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_usage_flags( &attributes, usage_flags );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, type );
 
     switch( generation_method )
     {
         case IMPORT_KEY:
             /* Import the key */
-            PSA_ASSERT( psa_import_key_to_handle( handle, type,
+            PSA_ASSERT( psa_import_key( &attributes, &handle,
                                         data->x, data->len ) );
             break;
 
         case GENERATE_KEY:
             /* Generate a key */
-            PSA_ASSERT( psa_generate_key_to_handle( handle, type, bits,
-                                          NULL, 0 ) );
+            PSA_ASSERT( psa_generate_key( &attributes, &handle,
+                                          bits, NULL, 0 ) );
             break;
 
         case DERIVE_KEY:
-            /* Create base key */
-            PSA_ASSERT( psa_allocate_key( &base_key ) );
-            psa_key_policy_set_usage( &base_policy_set, PSA_KEY_USAGE_DERIVE,
-                                      base_policy_alg );
-            PSA_ASSERT( psa_set_key_policy(
-                            base_key, &base_policy_set ) );
-            PSA_ASSERT( psa_import_key_to_handle( base_key, PSA_KEY_TYPE_DERIVE,
-                                        data->x, data->len ) );
-            /* Derive a key. */
-            PSA_ASSERT( psa_key_derivation( &generator, base_key,
-                                            base_policy_alg,
-                                            NULL, 0, NULL, 0,
-                                            export_size ) );
-            PSA_ASSERT( psa_generator_import_key_to_handle(
-                            handle, PSA_KEY_TYPE_RAW_DATA,
-                            bits, &generator ) );
+            {
+                /* Create base key */
+                psa_algorithm_t derive_alg = PSA_ALG_HKDF( PSA_ALG_SHA_256 );
+                psa_key_attributes_t base_attributes = PSA_KEY_ATTRIBUTES_INIT;
+                psa_set_key_usage_flags( &base_attributes,
+                                         PSA_KEY_USAGE_DERIVE );
+                psa_set_key_algorithm( &base_attributes, derive_alg );
+                psa_set_key_type( &base_attributes, PSA_KEY_TYPE_DERIVE );
+                PSA_ASSERT( psa_import_key( &base_attributes, &base_key,
+                                            data->x, data->len ) );
+                /* Derive a key. */
+                PSA_ASSERT( psa_key_derivation_setup( &generator, derive_alg ) );
+                PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                          PSA_KDF_STEP_SECRET,
+                                                          base_key ) );
+                PSA_ASSERT( psa_key_derivation_input_bytes(
+                                &generator, PSA_KDF_STEP_INFO,
+                                NULL, 0 ) );
+                PSA_ASSERT( psa_generator_import_key( &attributes, &handle,
+                                                      bits, &generator ) );
+                PSA_ASSERT( psa_generator_abort( &generator ) );
+                PSA_ASSERT( psa_destroy_key( base_key ) );
+                base_key = 0;
+            }
             break;
     }
+    psa_reset_key_attributes( &attributes );
 
-    /* Export the key */
-    TEST_EQUAL( psa_export_key( handle,
-                                first_export, export_size,
-                                &first_exported_length ),
-                export_status );
+    /* Export the key if permitted by the key policy. */
+    if( usage_flags & PSA_KEY_USAGE_EXPORT )
+    {
+        PSA_ASSERT( psa_export_key( handle,
+                                    first_export, export_size,
+                                    &first_exported_length ) );
+        if( generation_method == IMPORT_KEY )
+            ASSERT_COMPARE( data->x, data->len,
+                            first_export, first_exported_length );
+    }
 
     /* Shutdown and restart */
     mbedtls_psa_crypto_free();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Check key slot still contains key data */
-    PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, 1,
+    PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
                               &handle ) );
-    PSA_ASSERT( psa_get_key_information(
-                    handle, &type_get, &bits_get ) );
-    TEST_EQUAL( type_get, type );
-    TEST_EQUAL( bits_get, (size_t) bits );
+    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
+    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_EQUAL( psa_get_key_lifetime( &attributes ),
+                PSA_KEY_LIFETIME_PERSISTENT );
+    TEST_EQUAL( psa_get_key_type( &attributes ), type );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), bits );
+    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
+    TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
 
-    PSA_ASSERT( psa_get_key_policy( handle, &policy_get ) );
-    TEST_EQUAL( psa_key_policy_get_usage( &policy_get ), policy_usage );
-    TEST_EQUAL( psa_key_policy_get_algorithm( &policy_get ), policy_alg );
-
-    /* Export the key again */
-    TEST_EQUAL( psa_export_key( handle,
-                                second_export, export_size,
-                                &second_exported_length ),
-                export_status );
-
-    if( export_status == PSA_SUCCESS )
+    /* Export the key again if permitted by the key policy. */
+    if( usage_flags & PSA_KEY_USAGE_EXPORT )
     {
+        PSA_ASSERT( psa_export_key( handle,
+                                    second_export, export_size,
+                                    &second_exported_length ) );
         ASSERT_COMPARE( first_export, first_exported_length,
                         second_export, second_exported_length );
-
-        switch( generation_method )
-        {
-            case IMPORT_KEY:
-                ASSERT_COMPARE( data->x, data->len,
-                                first_export, first_exported_length );
-                break;
-            default:
-                break;
-        }
     }
 
     /* Do something with the key according to its type and permitted usage. */
-    if( ! exercise_key( handle, policy_usage, policy_alg ) )
+    if( ! exercise_key( handle, usage_flags, alg ) )
         goto exit;
 
 exit:
     mbedtls_free( first_export );
     mbedtls_free( second_export );
+    psa_generator_abort( &generator );
+    psa_destroy_key( base_key );
+    if( handle == 0 )
+    {
+        /* In case there was a test failure after creating the persistent key
+         * but while it was not open, try to re-open the persistent key
+         * to delete it. */
+        psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id, &handle );
+    }
     psa_destroy_key( handle );
     mbedtls_psa_crypto_free();
 }
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.data b/tests/suites/test_suite_psa_crypto_persistent_key.data
index f97a5e0..e8927b8 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.data
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.data
@@ -26,11 +26,7 @@
 
 Persistent key destroy
 depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C
-persistent_key_destroy:1:1:PSA_KEY_TYPE_RSA_KEYPAIR:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_TYPE_RAW_DATA:"deadbeef"
-
-Persistent key destroy missing key
-depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C
-persistent_key_destroy:1:0:PSA_KEY_TYPE_RSA_KEYPAIR:"":PSA_KEY_TYPE_RAW_DATA:"deadbeef"
+persistent_key_destroy:1:0:PSA_KEY_TYPE_RSA_KEYPAIR:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_TYPE_RAW_DATA:"deadbeef"
 
 Persistent key import
 depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.function b/tests/suites/test_suite_psa_crypto_persistent_key.function
index 245eeef..537fe93 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.function
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.function
@@ -87,6 +87,7 @@
     psa_key_handle_t handle = 0;
     uint8_t *data = NULL;
     size_t data_length = PSA_CRYPTO_MAX_STORAGE_SIZE;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     if( data_too_large )
         data_length += 1;
@@ -95,10 +96,10 @@
 
     PSA_ASSERT( psa_crypto_init() );
 
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
-                                &handle ) );
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
 
-    TEST_EQUAL( psa_import_key_to_handle( handle, PSA_KEY_TYPE_RAW_DATA,
+    TEST_EQUAL( psa_import_key( &attributes, &handle,
                                 data, data_length ),
                 expected_status );
 
@@ -110,7 +111,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void persistent_key_destroy( int key_id_arg, int should_store,
+void persistent_key_destroy( int key_id_arg, int restart,
                              int first_type_arg, data_t *first_data,
                              int second_type_arg, data_t *second_data )
 {
@@ -118,18 +119,25 @@
     psa_key_handle_t handle = 0;
     psa_key_type_t first_type = (psa_key_type_t) first_type_arg;
     psa_key_type_t second_type = (psa_key_type_t) second_type_arg;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     PSA_ASSERT( psa_crypto_init() );
 
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
-                                &handle ) );
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_type( &attributes, first_type );
 
-    if( should_store == 1 )
+    PSA_ASSERT( psa_import_key( &attributes, &handle,
+                                first_data->x, first_data->len ) );
+
+    if( restart )
     {
-        PSA_ASSERT( psa_import_key_to_handle(
-                        handle, first_type,
-                        first_data->x, first_data->len ) );
+        psa_close_key( handle );
+        mbedtls_psa_crypto_free();
+        PSA_ASSERT( psa_crypto_init() );
+        PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
+                                  &handle ) );
     }
+    TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 1 );
 
     /* Destroy the key */
     PSA_ASSERT( psa_destroy_key( handle ) );
@@ -145,11 +153,10 @@
     PSA_ASSERT( psa_crypto_init() );
 
     /* Create another key in the same slot */
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
-                                &handle ) );
-    PSA_ASSERT( psa_import_key_to_handle(
-                    handle, second_type,
-                    second_data->x, second_data->len ) );
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_type( &attributes, second_type );
+    PSA_ASSERT( psa_import_key( &attributes, &handle,
+                                second_data->x, second_data->len ) );
 
 exit:
     mbedtls_psa_crypto_free();
@@ -161,16 +168,16 @@
 void persistent_key_import( int key_id_arg, int type_arg, data_t *data,
                             int expected_status )
 {
-    psa_key_lifetime_t lifetime;
     psa_key_id_t key_id = (psa_key_id_t) key_id_arg;
     psa_key_type_t type = (psa_key_type_t) type_arg;
     psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     PSA_ASSERT( psa_crypto_init() );
 
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
-                                &handle ) );
-    TEST_EQUAL( psa_import_key_to_handle( handle, type, data->x, data->len ),
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_type( &attributes, type );
+    TEST_EQUAL( psa_import_key( &attributes, &handle, data->x, data->len ),
                 expected_status );
 
     if( expected_status != PSA_SUCCESS )
@@ -179,8 +186,14 @@
         goto exit;
     }
 
-    PSA_ASSERT( psa_get_key_lifetime_from_handle( handle, &lifetime ) );
-    TEST_EQUAL( lifetime, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_reset_key_attributes( &attributes );
+    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
+    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_EQUAL( psa_get_key_lifetime( &attributes ),
+                PSA_KEY_LIFETIME_PERSISTENT );
+    TEST_EQUAL( psa_get_key_type( &attributes ), type );
+    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
+    TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
 
 exit:
     psa_destroy_persistent_key( key_id );
@@ -198,34 +211,30 @@
     unsigned char *exported = NULL;
     size_t export_size = data->len;
     size_t exported_length;
-    psa_key_type_t got_type;
-    size_t got_bits;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
-    psa_key_lifetime_t lifetime_get;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     ASSERT_ALLOC( exported, export_size );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_create_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
-                                &handle ) );
-
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_EXPORT,
-                              PSA_ALG_VENDOR_FLAG );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
+    psa_set_key_type( &attributes, type );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
 
     /* Import the key */
-    PSA_ASSERT( psa_import_key_to_handle( handle, type,
+    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                 data->x, data->len ) );
 
-    PSA_ASSERT( psa_get_key_lifetime_from_handle( handle, &lifetime_get ) );
-    TEST_EQUAL( lifetime_get, PSA_KEY_LIFETIME_PERSISTENT );
-
     /* Test the key information */
-    PSA_ASSERT( psa_get_key_information(
-                    handle, &got_type, &got_bits ) );
-    TEST_EQUAL( got_type, type );
-    TEST_EQUAL( got_bits, (size_t) expected_bits );
+    psa_reset_key_attributes( &attributes );
+    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
+    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_EQUAL( psa_get_key_lifetime( &attributes ),
+                PSA_KEY_LIFETIME_PERSISTENT );
+    TEST_EQUAL( psa_get_key_type( &attributes ), type );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), (size_t) expected_bits );
+    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), PSA_KEY_USAGE_EXPORT );
+    TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
 
     TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 1 );