Update SE support to pass a location when registering a driver

Now that lifetimes have structures and secure element drivers handle
all the lifetimes with a certain location, update driver registration
to take a location as argument rather than a lifetime.

This commit updates the tests.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index 869fa3f..0c28579 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -1380,7 +1380,7 @@
  * \return PSA_ERROR_NOT_PERMITTED
  */
 psa_status_t psa_register_se_driver(
-    psa_key_lifetime_t lifetime,
+    psa_key_location_t location,
     const psa_drv_se_t *methods);
 
 /**@}*/
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.data b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
index 0c2411b..55c3426 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -6,14 +6,11 @@
 Register SE driver: good, again
 register_one:2:PSA_DRV_SE_HAL_VERSION:PSA_SUCCESS
 
-Register SE driver: invalid lifetime (0)
+Register SE driver: invalid location (0)
 register_one:0:PSA_DRV_SE_HAL_VERSION:PSA_ERROR_INVALID_ARGUMENT
 
-Register SE driver: invalid lifetime (VOLATILE)
-register_one:PSA_KEY_LIFETIME_VOLATILE:PSA_DRV_SE_HAL_VERSION:PSA_ERROR_INVALID_ARGUMENT
-
-Register SE driver: invalid lifetime (PERSISTENT)
-register_one:PSA_KEY_LIFETIME_PERSISTENT:PSA_DRV_SE_HAL_VERSION:PSA_ERROR_INVALID_ARGUMENT
+Register SE driver: invalid location (local)
+register_one:PSA_KEY_LOCATION_LOCAL_STORAGE:PSA_DRV_SE_HAL_VERSION:PSA_ERROR_INVALID_ARGUMENT
 
 Register SE driver: invalid version (ancient)
 register_one:2:0x00000003:PSA_ERROR_NOT_SUPPORTED
@@ -121,7 +118,7 @@
 generate_key_smoke:PSA_KEY_TYPE_HMAC:256:PSA_ALG_HMAC( PSA_ALG_SHA_256 )
 
 Key registration: smoke test
-register_key_smoke_test:MIN_DRIVER_LIFETIME:1:PSA_SUCCESS
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:1:PSA_SUCCESS
 
 Key registration: invalid lifetime (volatile)
 register_key_smoke_test:PSA_KEY_LIFETIME_VOLATILE:1:PSA_ERROR_INVALID_ARGUMENT
@@ -130,13 +127,13 @@
 register_key_smoke_test:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_ERROR_INVALID_ARGUMENT
 
 Key registration: invalid lifetime (no registered driver)
-register_key_smoke_test:MIN_DRIVER_LIFETIME + 1:1:PSA_ERROR_INVALID_ARGUMENT
+register_key_smoke_test:PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION + 1 ):1:PSA_ERROR_INVALID_ARGUMENT
 
 Key registration: rejected
-register_key_smoke_test:MIN_DRIVER_LIFETIME:0:PSA_ERROR_NOT_PERMITTED
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:0:PSA_ERROR_NOT_PERMITTED
 
 Key registration: not supported
-register_key_smoke_test:MIN_DRIVER_LIFETIME:-1:PSA_ERROR_NOT_SUPPORTED
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:-1:PSA_ERROR_NOT_SUPPORTED
 
 Import-sign-verify: sign in driver, ECDSA
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index b468d5e..f95f7e5 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -18,8 +18,14 @@
 /* Test driver helpers */
 /****************************************************************/
 
-/** The minimum valid lifetime value for a secure element driver. */
-#define MIN_DRIVER_LIFETIME 2
+/** The minimum valid location value for a secure element driver. */
+#define MIN_DRIVER_LOCATION 1
+
+/** The location and lifetime used for tests that use a single driver. */
+#define TEST_DRIVER_LOCATION 1
+#define TEST_SE_PERSISTENT_LIFETIME                             \
+    ( PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(           \
+        PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION ) )
 
 /** The driver detected a condition that shouldn't happen.
  * This is probably a bug in the library. */
@@ -547,7 +553,7 @@
             psa_get_key_lifetime( &actual_attributes );
         psa_status_t status = psa_get_key_slot_number( &actual_attributes,
                                                        &actual_slot_number );
-        if( lifetime < MIN_DRIVER_LIFETIME )
+        if( PSA_KEY_LIFETIME_GET_LOCATION( lifetime ) < MIN_DRIVER_LOCATION )
         {
             /* The key is not in a secure element. */
             TEST_EQUAL( status, PSA_ERROR_INVALID_ARGUMENT );
@@ -571,23 +577,23 @@
     return( ok );
 }
 
-/* Get the file UID corresponding to the specified lifetime.
+/* Get the file UID corresponding to the specified location.
  * If this changes, the storage format version must change.
  * See psa_get_se_driver_its_file_uid() in psa_crypto_se.c.
  */
-psa_storage_uid_t file_uid_for_lifetime( psa_key_lifetime_t lifetime )
+psa_storage_uid_t file_uid_for_location( psa_key_location_t location )
 {
-    if( lifetime > PSA_MAX_SE_LIFETIME )
+    if( location > PSA_MAX_SE_LOCATION )
         return( 0 );
-    return( 0xfffffe00 + lifetime );
+    return( 0xfffffe00 + location );
 }
 
 /* Check that the persistent data of a driver has its expected content. */
-static int check_persistent_data( psa_key_lifetime_t lifetime,
+static int check_persistent_data( psa_key_location_t location,
                                   const void *expected_data,
                                   size_t size )
 {
-    psa_storage_uid_t uid = file_uid_for_lifetime( lifetime );
+    psa_storage_uid_t uid = file_uid_for_location( location );
     struct psa_storage_info_t info;
     uint8_t *loaded = NULL;
     int ok = 0;
@@ -737,7 +743,7 @@
 static void psa_purge_storage( void )
 {
     psa_key_id_t id;
-    psa_key_lifetime_t lifetime;
+    psa_key_location_t location;
     /* The tests may have potentially created key ids from 1 to
      * MAX_KEY_ID_FOR_TEST. In addition, run the destroy function on key id
      * 0, which file-based storage uses as a temporary file. */
@@ -746,8 +752,8 @@
     /* Purge the transaction file. */
     psa_crypto_stop_transaction( );
     /* Purge driver persistent data. */
-    for( lifetime = 0; lifetime < PSA_MAX_SE_LIFETIME; lifetime++ )
-        psa_destroy_se_persistent_data( lifetime );
+    for( location = 0; location < PSA_MAX_SE_LOCATION; location++ )
+        psa_destroy_se_persistent_data( location );
 }
 
 /* END_HEADER */
@@ -758,7 +764,7 @@
  */
 
 /* BEGIN_CASE */
-void register_one( int lifetime, int version, int expected_status_arg )
+void register_one( int location, int version, int expected_status_arg )
 {
     psa_status_t expected_status = expected_status_arg;
     psa_drv_se_t driver;
@@ -766,7 +772,7 @@
     memset( &driver, 0, sizeof( driver ) );
     driver.hal_version = version;
 
-    TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
+    TEST_EQUAL( psa_register_se_driver( location, &driver ),
                 expected_status );
 
     PSA_ASSERT( psa_crypto_init( ) );
@@ -780,16 +786,16 @@
 void register_twice( int count )
 {
     psa_drv_se_t driver;
-    psa_key_lifetime_t lifetime;
-    psa_key_lifetime_t max = MIN_DRIVER_LIFETIME + count;
+    psa_key_location_t location;
+    psa_key_location_t max = MIN_DRIVER_LOCATION + count;
 
     memset( &driver, 0, sizeof( driver ) );
     driver.hal_version = PSA_DRV_SE_HAL_VERSION;
 
-    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
-        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
-    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
-        TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
+    for( location = MIN_DRIVER_LOCATION; location < max; location++ )
+        PSA_ASSERT( psa_register_se_driver( location, &driver ) );
+    for( location = MIN_DRIVER_LOCATION; location < max; location++ )
+        TEST_EQUAL( psa_register_se_driver( location, &driver ),
                     PSA_ERROR_ALREADY_EXISTS );
 
     PSA_ASSERT( psa_crypto_init( ) );
@@ -803,16 +809,16 @@
 void register_max( )
 {
     psa_drv_se_t driver;
-    psa_key_lifetime_t lifetime;
-    psa_key_lifetime_t max = MIN_DRIVER_LIFETIME + PSA_MAX_SE_DRIVERS;
+    psa_key_location_t location;
+    psa_key_location_t max = MIN_DRIVER_LOCATION + PSA_MAX_SE_DRIVERS;
 
     memset( &driver, 0, sizeof( driver ) );
     driver.hal_version = PSA_DRV_SE_HAL_VERSION;
 
-    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
-        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    for( location = MIN_DRIVER_LOCATION; location < max; location++ )
+        PSA_ASSERT( psa_register_se_driver( location, &driver ) );
 
-    TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
+    TEST_EQUAL( psa_register_se_driver( location, &driver ),
                 PSA_ERROR_INSUFFICIENT_MEMORY );
 
     PSA_ASSERT( psa_crypto_init( ) );
@@ -827,7 +833,8 @@
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -846,7 +853,7 @@
     key_management.p_export = ram_export;
     ram_min_slot = min_slot;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Create a key. */
@@ -857,7 +864,7 @@
     PSA_ASSERT( psa_import_key( &attributes,
                                 key_material, sizeof( key_material ),
                                 &handle ) );
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &ram_shadow_slot_usage,
                                  sizeof( ram_shadow_slot_usage ) ) )
         goto exit;
@@ -866,9 +873,9 @@
     if( restart )
     {
         mbedtls_psa_crypto_free( );
-        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+        PSA_ASSERT( psa_register_se_driver( location, &driver ) );
         PSA_ASSERT( psa_crypto_init( ) );
-        if( ! check_persistent_data( lifetime,
+        if( ! check_persistent_data( location,
                                      &ram_shadow_slot_usage,
                                      sizeof( ram_shadow_slot_usage ) ) )
             goto exit;
@@ -894,7 +901,7 @@
 
     PSA_ASSERT( psa_destroy_key( handle ) );
     handle = 0;
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &ram_shadow_slot_usage,
                                  sizeof( ram_shadow_slot_usage ) ) )
         goto exit;
@@ -921,7 +928,8 @@
     psa_status_t status;
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -937,7 +945,7 @@
     key_management.p_destroy = ram_destroy;
     key_management.p_export = ram_export;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Create a key. */
@@ -953,7 +961,7 @@
 
     if( status != PSA_SUCCESS )
         goto exit;
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &ram_shadow_slot_usage,
                                  sizeof( ram_shadow_slot_usage ) ) )
         goto exit;
@@ -962,9 +970,9 @@
     if( restart )
     {
         mbedtls_psa_crypto_free( );
-        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+        PSA_ASSERT( psa_register_se_driver( location, &driver ) );
         PSA_ASSERT( psa_crypto_init( ) );
-        if( ! check_persistent_data( lifetime,
+        if( ! check_persistent_data( location,
                                      &ram_shadow_slot_usage,
                                      sizeof( ram_shadow_slot_usage ) ) )
             goto exit;
@@ -980,7 +988,7 @@
 
     PSA_ASSERT( psa_destroy_key( handle ) );
     handle = 0;
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &ram_shadow_slot_usage,
                                  sizeof( ram_shadow_slot_usage ) ) )
         goto exit;
@@ -1002,7 +1010,8 @@
     psa_algorithm_t alg = alg_arg;
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -1016,7 +1025,7 @@
     key_management.p_import = null_import;
     key_management.p_destroy = null_destroy;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Create a key. */
@@ -1031,7 +1040,7 @@
     PSA_ASSERT( psa_import_key( &attributes,
                                 key_material->x, key_material->len,
                                 &handle ) );
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
 
@@ -1041,9 +1050,9 @@
 
     /* Restart and try again. */
     mbedtls_psa_crypto_free( );
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
     PSA_ASSERT( psa_open_key( id, &handle ) );
@@ -1053,7 +1062,7 @@
     /* We're done. */
     PSA_ASSERT( psa_destroy_key( handle ) );
     handle = 0;
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
     TEST_EQUAL( psa_open_key( id, &handle ),
@@ -1073,7 +1082,8 @@
     size_t bits = bits_arg;
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -1086,7 +1096,7 @@
     key_management.p_allocate = counter_allocate;
     /* No p_generate method */
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -1111,7 +1121,8 @@
     psa_algorithm_t alg = alg_arg;
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -1125,7 +1136,7 @@
     key_management.p_generate = null_generate;
     key_management.p_destroy = null_destroy;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Create a key. */
@@ -1139,7 +1150,7 @@
     psa_set_key_type( &attributes, type );
     psa_set_key_bits( &attributes, bits );
     PSA_ASSERT( psa_generate_key( &attributes, &handle ) );
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
 
@@ -1149,9 +1160,9 @@
 
     /* Restart and try again. */
     mbedtls_psa_crypto_free( );
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
     PSA_ASSERT( psa_open_key( id, &handle ) );
@@ -1161,7 +1172,7 @@
     /* We're done. */
     PSA_ASSERT( psa_destroy_key( handle ) );
     handle = 0;
-    if( ! check_persistent_data( lifetime,
+    if( ! check_persistent_data( location,
                                  &shadow_counter, sizeof( shadow_counter ) ) )
         goto exit;
     TEST_EQUAL( psa_open_key( id, &handle ),
@@ -1190,7 +1201,8 @@
     psa_drv_se_key_management_t key_management;
     psa_drv_se_asymmetric_t asymmetric;
 
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t drv_handle = 0; /* key managed by the driver */
     psa_key_handle_t sw_handle = 0; /* transparent key */
@@ -1229,7 +1241,7 @@
     }
     asymmetric.p_verify = ram_verify;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Prepare to create two keys with the same key material: a transparent
@@ -1347,6 +1359,7 @@
                               int expected_status_arg )
 {
     psa_key_lifetime_t lifetime = lifetime_arg;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_status_t expected_status = expected_status_arg;
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
@@ -1371,7 +1384,7 @@
             ( validate > 0 ? PSA_SUCCESS : PSA_ERROR_NOT_PERMITTED );
     }
 
-    PSA_ASSERT( psa_register_se_driver( MIN_DRIVER_LIFETIME, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( MIN_DRIVER_LOCATION, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -1395,7 +1408,7 @@
 
     /* Restart and try again. */
     PSA_DONE( );
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
     PSA_ASSERT( psa_open_key( id, &handle ) );
     if( ! check_key_attributes( handle, &attributes ) )
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
index f60bd76..23e035a 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
@@ -1,11 +1,14 @@
 SE init mock test: success
-mock_init:2:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:1
+mock_init:1:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:1
 
 SE init mock test: failure
-mock_init:2:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE:1
+mock_init:1:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE:1
 
-SE init mock test: invalid lifetime
-mock_init:1:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE:PSA_SUCCESS:0
+SE init mock test: invalid location (0)
+mock_init:0:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE:PSA_SUCCESS:0
+
+SE init mock test: location not supported (INT_MAX)
+mock_init:INT_MAX:PSA_ERROR_NOT_SUPPORTED:PSA_ERROR_BAD_STATE:PSA_SUCCESS:0
 
 SE key importing mock test
 mock_import:PSA_SUCCESS:PSA_SUCCESS:0:PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
index 9f17b84..f6acb07 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
@@ -5,10 +5,16 @@
 #include "psa_crypto_se.h"
 #include "psa_crypto_storage.h"
 
+/** The location and lifetime used for tests that use a single driver. */
+#define TEST_DRIVER_LOCATION 1
+#define TEST_SE_PERSISTENT_LIFETIME                             \
+    ( PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(           \
+        PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION ) )
+
 static struct
 {
     uint16_t called;
-    psa_key_lifetime_t lifetime;
+    psa_key_location_t location;
     psa_status_t return_value;
 } mock_init_data;
 
@@ -84,7 +90,7 @@
 static void psa_purge_storage( void )
 {
     psa_key_id_t id;
-    psa_key_lifetime_t lifetime;
+    psa_key_location_t location;
     /* The tests may have potentially created key ids from 1 to
      * MAX_KEY_ID_FOR_TEST. In addition, run the destroy function on key id
      * 0, which file-based storage uses as a temporary file. */
@@ -93,8 +99,8 @@
     /* Purge the transaction file. */
     psa_crypto_stop_transaction( );
     /* Purge driver persistent data. */
-    for( lifetime = 0; lifetime < PSA_MAX_SE_LIFETIME; lifetime++ )
-        psa_destroy_se_persistent_data( lifetime );
+    for( location = 0; location < PSA_MAX_SE_LOCATION; location++ )
+        psa_destroy_se_persistent_data( location );
 }
 
 static void mock_teardown( void )
@@ -113,13 +119,13 @@
 
 static psa_status_t mock_init( psa_drv_se_context_t *drv_context,
                                void *persistent_data,
-                               psa_key_lifetime_t lifetime )
+                               psa_key_location_t location )
 {
     (void) drv_context;
     (void) persistent_data;
 
     mock_init_data.called++;
-    mock_init_data.lifetime = lifetime;
+    mock_init_data.location = location;
     return( mock_init_data.return_value );
 }
 
@@ -279,13 +285,13 @@
  */
 
 /* BEGIN_CASE */
-void mock_init( int lifetime_arg,
+void mock_init( int location_arg,
                 int expected_register_status_arg,
                 int driver_status_arg,
                 int expected_psa_status_arg,
                 int expected_called )
 {
-    psa_key_lifetime_t lifetime = lifetime_arg;
+    psa_key_location_t location = location_arg;
     psa_status_t expected_register_status = expected_register_status_arg;
     psa_status_t driver_status = driver_status_arg;
     psa_status_t expected_psa_status = expected_psa_status_arg;
@@ -297,7 +303,7 @@
 
     mock_init_data.return_value = driver_status;
 
-    TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
+    TEST_EQUAL( psa_register_se_driver( location, &driver ),
                 expected_register_status );
 
     psa_crypto_init_called = 1;
@@ -305,7 +311,7 @@
 
     TEST_EQUAL( mock_init_data.called, expected_called );
     if( expected_called )
-        TEST_EQUAL( mock_init_data.lifetime, lifetime );
+        TEST_EQUAL( mock_init_data.location, location );
 
 exit:
     if( psa_crypto_init_called )
@@ -322,7 +328,8 @@
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -339,7 +346,7 @@
     key_management.p_destroy = mock_destroy;
     key_management.p_allocate = mock_allocate;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -378,7 +385,8 @@
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -397,7 +405,7 @@
     key_management.p_destroy = mock_destroy;
     key_management.p_allocate = mock_allocate;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -431,7 +439,8 @@
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -446,7 +455,7 @@
     key_management.p_destroy = mock_destroy;
     key_management.p_allocate = mock_allocate;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -485,7 +494,8 @@
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -503,7 +513,7 @@
     key_management.p_destroy = mock_destroy;
     key_management.p_allocate = mock_allocate;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -534,7 +544,8 @@
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
     psa_drv_se_asymmetric_t asymmetric;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -559,7 +570,7 @@
     driver.asymmetric = &asymmetric;
     asymmetric.p_sign = mock_sign;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -594,7 +605,8 @@
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
     psa_drv_se_asymmetric_t asymmetric;
-    psa_key_lifetime_t lifetime = 2;
+    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -618,7 +630,7 @@
     driver.asymmetric = &asymmetric;
     asymmetric.p_verify = mock_verify;
 
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );