diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index ba1e04a..48c33d7 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -38,9 +38,13 @@
     /* Count the amount of times one of the key management driver functions
      * is called. */
     unsigned long hits;
+    /* Record the source of the function call. */
+    psa_key_location_t source;
 } mbedtls_test_driver_key_management_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0 }
+/* 0x800000 is a vendor-specific location, unused by the PSA, overwritten
+ * in tests that expect a different value. */
+#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0x800000 }
 static inline mbedtls_test_driver_key_management_hooks_t
     mbedtls_test_driver_key_management_hooks_init( void )
 {
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index 5028073..3cfd595 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -29,6 +29,8 @@
 #include "mbedtls/error.h"
 
 #include "test/drivers/key_management.h"
+#include "test/drivers/test_driver.h"
+
 #include "test/random.h"
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
@@ -260,6 +262,7 @@
     size_t *bits)
 {
     ++mbedtls_test_driver_key_management_hooks.hits;
+    mbedtls_test_driver_key_management_hooks.source = PSA_KEY_LOCATION_LOCAL_STORAGE;
 
     if( mbedtls_test_driver_key_management_hooks.forced_status != PSA_SUCCESS )
         return( mbedtls_test_driver_key_management_hooks.forced_status );
@@ -326,6 +329,12 @@
     size_t *key_buffer_length,
     size_t *bits)
 {
+    ++mbedtls_test_driver_key_management_hooks.hits;
+    mbedtls_test_driver_key_management_hooks.source = PSA_CRYPTO_TEST_DRIVER_LOCATION;
+
+    if( mbedtls_test_driver_key_management_hooks.forced_status != PSA_SUCCESS )
+        return( mbedtls_test_driver_key_management_hooks.forced_status );
+
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_type_t type = psa_get_key_type( attributes );
     /* This buffer will be used as an intermediate placeholder for
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index ea6c9b3..ab25726 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -83,23 +83,35 @@
 
 validate key through transparent driver: good private key
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
-validate_key:PSA_SUCCESS:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_SUCCESS
+validate_key:PSA_SUCCESS:PSA_KEY_LOCATION_LOCAL_STORAGE:130:1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_SUCCESS
 
 validate key through transparent driver: good public key
 depends_on:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
-validate_key:PSA_SUCCESS:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_SUCCESS
+validate_key:PSA_SUCCESS:PSA_KEY_LOCATION_LOCAL_STORAGE:131:1:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_SUCCESS
 
 validate key through transparent driver: fallback private key
 depends_on:MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
-validate_key:PSA_ERROR_NOT_SUPPORTED:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_SUCCESS
+validate_key:PSA_ERROR_NOT_SUPPORTED:PSA_KEY_LOCATION_LOCAL_STORAGE:132:1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_SUCCESS
 
 validate key through transparent driver: fallback public key
 depends_on:MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
-validate_key:PSA_ERROR_NOT_SUPPORTED:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_SUCCESS
+validate_key:PSA_ERROR_NOT_SUPPORTED:PSA_KEY_LOCATION_LOCAL_STORAGE:133:1:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_SUCCESS
 
 validate key through transparent driver: error
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR
-validate_key:PSA_ERROR_GENERIC_ERROR:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ERROR_GENERIC_ERROR
+validate_key:PSA_ERROR_GENERIC_ERROR:PSA_KEY_LOCATION_LOCAL_STORAGE:134:1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ERROR_GENERIC_ERROR
+
+validate key through opaque driver: good private key
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
+validate_key:PSA_SUCCESS:PSA_CRYPTO_TEST_DRIVER_LOCATION:130:1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_SUCCESS
+
+validate key through opaque driver: good public key
+depends_on:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
+validate_key:PSA_SUCCESS:PSA_CRYPTO_TEST_DRIVER_LOCATION:131:1:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04dea5e45d0ea37fc566232a508f4ad20ea13d47e4bf5fa4d54a57a0ba012042087097496efc583fed8b24a5b9be9a51de063f5a00a8b698a16fd7f29b5485f320":PSA_SUCCESS
+
+validate key through opaque driver: error
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR
+validate_key:PSA_ERROR_GENERIC_ERROR:PSA_CRYPTO_TEST_DRIVER_LOCATION:134:1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ERROR_GENERIC_ERROR
 
 export_key private to public through driver: fake
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 8b7f413..eab258d 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -314,10 +314,15 @@
 
 /* BEGIN_CASE */
 void validate_key( int force_status_arg,
+                   int location,
+                   int owner_id_arg,
+                   int id_arg,
                    int key_type_arg,
                    data_t *key_input,
                    int expected_status_arg )
 {
+    psa_key_lifetime_t lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_PERSISTENCE_DEFAULT, location);
+    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( owner_id_arg, id_arg );
     psa_status_t force_status = force_status_arg;
     psa_status_t expected_status = expected_status_arg;
     psa_key_type_t key_type = key_type_arg;
@@ -327,8 +332,10 @@
     mbedtls_test_driver_key_management_hooks =
         mbedtls_test_driver_key_management_hooks_init();
 
+    psa_set_key_id( &attributes, id );
     psa_set_key_type( &attributes,
                       key_type );
+    psa_set_key_lifetime( &attributes, lifetime );
     psa_set_key_bits( &attributes, 0 );
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
 
@@ -339,6 +346,7 @@
     actual_status = psa_import_key( &attributes, key_input->x, key_input->len, &key );
     TEST_EQUAL( mbedtls_test_driver_key_management_hooks.hits, 1 );
     TEST_EQUAL( actual_status, expected_status );
+    TEST_EQUAL( mbedtls_test_driver_key_management_hooks.source, location );
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( key );