diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 7eebfcf..413df0a 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4120,6 +4120,59 @@
 }
 #endif /* MBEDTLS_DES_C */
 
+static psa_status_t psa_generator_import_key_internal(
+    psa_key_slot_t *slot,
+    size_t bits,
+    psa_crypto_generator_t *generator )
+{
+    uint8_t *data = NULL;
+    size_t bytes = PSA_BITS_TO_BYTES( bits );
+    psa_status_t status;
+
+    if( ! key_type_is_raw_bytes( slot->type ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    if( bits % 8 != 0 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    data = mbedtls_calloc( 1, bytes );
+    if( data == NULL )
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
+    status = psa_generator_read( generator, data, bytes );
+    if( status != PSA_SUCCESS )
+        goto exit;
+#if defined(MBEDTLS_DES_C)
+    if( slot->type == PSA_KEY_TYPE_DES )
+        psa_des_set_key_parity( data, bytes );
+#endif /* MBEDTLS_DES_C */
+    status = psa_import_key_into_slot( slot, data, bytes );
+
+exit:
+    mbedtls_free( data );
+    return( status );
+}
+
+psa_status_t psa_generator_import_key( const psa_key_attributes_t *attributes,
+                                       psa_key_handle_t *handle,
+                                       size_t bits,
+                                       psa_crypto_generator_t *generator )
+{
+    psa_status_t status;
+    psa_key_slot_t *slot = NULL;
+    status = psa_start_key_creation( attributes, handle, &slot );
+    if( status == PSA_SUCCESS )
+    {
+        status = psa_generator_import_key_internal( slot, bits, generator );
+    }
+    if( status == PSA_SUCCESS )
+        status = psa_finish_key_creation( slot );
+    if( status != PSA_SUCCESS )
+    {
+        psa_fail_key_creation( slot );
+        *handle = 0;
+    }
+    return( status );
+}
+
 psa_status_t psa_generator_import_key_to_handle( psa_key_handle_t handle,
                                        psa_key_type_t type,
                                        size_t bits,
@@ -4873,24 +4926,19 @@
 }
 #endif /* MBEDTLS_PSA_INJECT_ENTROPY */
 
-psa_status_t psa_generate_key_to_handle( psa_key_handle_t handle,
-                               psa_key_type_t type,
-                               size_t bits,
-                               const void *extra,
-                               size_t extra_size )
+static psa_status_t psa_generate_key_internal( psa_key_slot_t *slot,
+                                               size_t bits,
+                                               const void *extra,
+                                               size_t extra_size )
 {
-    psa_key_slot_t *slot;
-    psa_status_t status;
+    psa_key_type_t type = slot->type;
 
     if( extra == NULL && extra_size != 0 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    status = psa_get_empty_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
     if( key_type_is_raw_bytes( type ) )
     {
+        psa_status_t status;
         status = prepare_raw_data_slot( type, bits, &slot->data.raw );
         if( status != PSA_SUCCESS )
             return( status );
@@ -4989,7 +5037,26 @@
 
         return( PSA_ERROR_NOT_SUPPORTED );
 
+    return( PSA_SUCCESS );
+}
+
+psa_status_t psa_generate_key_to_handle( psa_key_handle_t handle,
+                               psa_key_type_t type,
+                               size_t bits,
+                               const void *extra,
+                               size_t extra_size )
+{
+    psa_key_slot_t *slot;
+    psa_status_t status;
+
+    status = psa_get_empty_key_slot( handle, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
+
     slot->type = type;
+    status = psa_generate_key_internal( slot, bits, extra, extra_size );
+    if( status != PSA_SUCCESS )
+        slot->type = 0;
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
@@ -5001,6 +5068,30 @@
     return( status );
 }
 
+psa_status_t psa_generate_key( const psa_key_attributes_t *attributes,
+                               psa_key_handle_t *handle,
+                               size_t bits,
+                               const void *extra,
+                               size_t extra_size )
+{
+    psa_status_t status;
+    psa_key_slot_t *slot = NULL;
+    status = psa_start_key_creation( attributes, handle, &slot );
+    if( status == PSA_SUCCESS )
+    {
+        status = psa_generate_key_internal( slot, bits, extra, extra_size );
+    }
+    if( status == PSA_SUCCESS )
+        status = psa_finish_key_creation( slot );
+    if( status != PSA_SUCCESS )
+    {
+        psa_fail_key_creation( slot );
+        *handle = 0;
+    }
+    return( status );
+}
+
+
 
 /****************************************************************/
 /* Module setup */
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b9f0d5f..03ec2b0 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4320,7 +4320,7 @@
     uint8_t *output_buffer = NULL;
     size_t expected_capacity;
     size_t current_capacity;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status;
     unsigned i;
 
@@ -4334,11 +4334,11 @@
     ASSERT_ALLOC( output_buffer, output_buffer_size );
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
 
-    PSA_ASSERT( psa_import_key_to_handle( handle, PSA_KEY_TYPE_DERIVE,
+    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                 key_data->x,
                                 key_data->len ) );
 
@@ -4427,15 +4427,15 @@
     unsigned char output_buffer[16];
     size_t expected_capacity = requested_capacity;
     size_t current_capacity;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
 
-    PSA_ASSERT( psa_import_key_to_handle( handle, PSA_KEY_TYPE_DERIVE,
+    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                 key_data->x,
                                 key_data->len ) );
 
@@ -4514,16 +4514,16 @@
     psa_algorithm_t derived_alg = derived_alg_arg;
     size_t capacity = PSA_BITS_TO_BYTES( derived_bits );
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_type_t got_type;
     size_t got_bits;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &base_handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( base_handle, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( base_handle, PSA_KEY_TYPE_DERIVE,
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
+    PSA_ASSERT( psa_import_key( &attributes, &base_handle,
                                 key_data->x,
                                 key_data->len ) );
 
@@ -4532,11 +4532,10 @@
                                     salt->x, salt->len,
                                     label->x, label->len,
                                     capacity ) );
-    PSA_ASSERT( psa_allocate_key( &derived_handle ) );
-    psa_key_policy_set_usage( &policy, derived_usage, derived_alg );
-    PSA_ASSERT( psa_set_key_policy( derived_handle, &policy ) );
-    PSA_ASSERT( psa_generator_import_key_to_handle( derived_handle,
-                                          derived_type,
+    psa_set_key_usage_flags( &attributes, derived_usage );
+    psa_set_key_algorithm( &attributes, derived_alg );
+    psa_set_key_type( &attributes, derived_type );
+    PSA_ASSERT( psa_generator_import_key( &attributes, &derived_handle,
                                           derived_bits,
                                           &generator ) );
 
@@ -4577,17 +4576,18 @@
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
     uint8_t *output_buffer = NULL;
     uint8_t *export_buffer = NULL;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t base_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_attributes_t derived_attributes = PSA_KEY_ATTRIBUTES_INIT;
     size_t length;
 
     ASSERT_ALLOC( output_buffer, capacity );
     ASSERT_ALLOC( export_buffer, capacity );
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &base_handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( base_handle, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( base_handle, PSA_KEY_TYPE_DERIVE,
+    psa_set_key_usage_flags( &base_attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &base_attributes, alg );
+    psa_set_key_type( &base_attributes, PSA_KEY_TYPE_DERIVE );
+    PSA_ASSERT( psa_import_key( &base_attributes, &base_handle,
                                 key_data->x,
                                 key_data->len ) );
 
@@ -4606,11 +4606,10 @@
                                     salt->x, salt->len,
                                     label->x, label->len,
                                     capacity ) );
-    PSA_ASSERT( psa_allocate_key( &derived_handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_EXPORT, 0 );
-    PSA_ASSERT( psa_set_key_policy( derived_handle, &policy ) );
-    PSA_ASSERT( psa_generator_import_key_to_handle( derived_handle,
-                                          PSA_KEY_TYPE_RAW_DATA,
+    psa_set_key_usage_flags( &derived_attributes, PSA_KEY_USAGE_EXPORT );
+    psa_set_key_algorithm( &derived_attributes, 0 );
+    psa_set_key_type( &derived_attributes, PSA_KEY_TYPE_RAW_DATA );
+    PSA_ASSERT( psa_generator_import_key( &derived_attributes, &derived_handle,
                                           derived_bits,
                                           &generator ) );
     PSA_ASSERT( psa_export_key( derived_handle,
@@ -4618,10 +4617,7 @@
                                 &length ) );
     TEST_EQUAL( length, bytes1 );
     PSA_ASSERT( psa_destroy_key( derived_handle ) );
-    PSA_ASSERT( psa_allocate_key( &derived_handle ) );
-    PSA_ASSERT( psa_set_key_policy( derived_handle, &policy ) );
-    PSA_ASSERT( psa_generator_import_key_to_handle( derived_handle,
-                                          PSA_KEY_TYPE_RAW_DATA,
+    PSA_ASSERT( psa_generator_import_key( &derived_attributes, &derived_handle,
                                           PSA_BYTES_TO_BITS( bytes2 ),
                                           &generator ) );
     PSA_ASSERT( psa_export_key( derived_handle,
@@ -4653,16 +4649,16 @@
     psa_algorithm_t alg = alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t expected_status = expected_status_arg;
     psa_status_t status;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &our_key ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( our_key, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( our_key, our_key_type,
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, our_key_type );
+    PSA_ASSERT( psa_import_key( &attributes, &our_key,
                                 our_key_data->x,
                                 our_key_data->len ) );
 
@@ -4699,17 +4695,17 @@
     psa_key_handle_t our_key = 0;
     psa_algorithm_t alg = alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     unsigned char *output = NULL;
     size_t output_length = ~0;
 
     ASSERT_ALLOC( output, expected_output->len );
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &our_key ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( our_key, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( our_key, our_key_type,
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, our_key_type );
+    PSA_ASSERT( psa_import_key( &attributes, &our_key,
                                 our_key_data->x,
                                 our_key_data->len ) );
 
@@ -4737,16 +4733,16 @@
     psa_algorithm_t alg = alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     size_t actual_capacity;
     unsigned char output[16];
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &our_key ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( our_key, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( our_key, our_key_type,
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, our_key_type );
+    PSA_ASSERT( psa_import_key( &attributes, &our_key,
                                 our_key_data->x,
                                 our_key_data->len ) );
 
@@ -4796,7 +4792,7 @@
     psa_algorithm_t alg = alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     uint8_t *actual_output = NULL;
 
     ASSERT_ALLOC( actual_output, MAX( expected_output1->len,
@@ -4804,10 +4800,10 @@
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &our_key ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
-    PSA_ASSERT( psa_set_key_policy( our_key, &policy ) );
-    PSA_ASSERT( psa_import_key_to_handle( our_key, our_key_type,
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, our_key_type );
+    PSA_ASSERT( psa_import_key( &attributes, &our_key,
                                 our_key_data->x,
                                 our_key_data->len ) );
 
@@ -4913,23 +4909,23 @@
     size_t got_bits;
     psa_status_t expected_info_status =
         expected_status == PSA_SUCCESS ? PSA_SUCCESS : PSA_ERROR_DOES_NOT_EXIST;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, usage, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_set_key_usage_flags( &attributes, usage );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, type );
 
     /* Generate a key */
-    TEST_EQUAL( psa_generate_key_to_handle( handle, type, bits, NULL, 0 ),
+    TEST_EQUAL( psa_generate_key( &attributes, &handle, bits, NULL, 0 ),
                 expected_status );
+    if( expected_info_status != PSA_SUCCESS )
+        goto exit;
 
     /* Test the key information */
     TEST_EQUAL( psa_get_key_information( handle, &got_type, &got_bits ),
                 expected_info_status );
-    if( expected_info_status != PSA_SUCCESS )
-        goto exit;
     TEST_EQUAL( got_type, type );
     TEST_EQUAL( got_bits, bits );