diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 160076e..1910894 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -594,14 +594,10 @@
     {
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
             return( test_opaque_get_builtin_key(
                         slot_number,
                         attributes,
                         key_buffer, key_buffer_size, key_buffer_length ) );
-#else
-            return( PSA_ERROR_DOES_NOT_EXIST );
-#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
 #endif /* PSA_CRYPTO_DRIVER_TEST */
         default:
             (void) slot_number;
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index cf6fbb0..100fc18 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -29,14 +29,9 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
 #include <psa/crypto_driver_common.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
 #define PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT     0
 #define PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT   1
 
-extern const uint8_t test_driver_aes_key[16];
-extern const uint8_t test_driver_ecdsa_key[32];
-#endif
-
 typedef struct {
     /* If non-null, on success, copy this to the output. */
     void *forced_output;
diff --git a/tests/src/drivers/key_management.c b/tests/src/drivers/key_management.c
index ca00fe0..e908daf 100644
--- a/tests/src/drivers/key_management.c
+++ b/tests/src/drivers/key_management.c
@@ -41,7 +41,6 @@
 test_driver_key_management_hooks_t test_driver_key_management_hooks =
     TEST_DRIVER_KEY_MANAGEMENT_INIT;
 
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
 const uint8_t test_driver_aes_key[16] =
     { 0x36, 0x77, 0x39, 0x7A, 0x24, 0x43, 0x26, 0x46,
       0x29, 0x4A, 0x40, 0x4E, 0x63, 0x52, 0x66, 0x55 };
@@ -61,12 +60,6 @@
       0xbc, 0x25, 0x16, 0xc3, 0xd2, 0x70, 0x2d, 0x79,
       0x2f, 0x13, 0x1a, 0x92, 0x20, 0x95, 0xfd, 0x6c };
 
-static const psa_drv_slot_number_t aes_slot =
-    PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT;
-static const psa_drv_slot_number_t ecdsa_slot =
-    PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT;
-#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
-
 psa_status_t test_transparent_generate_key(
     const psa_key_attributes_t *attributes,
     uint8_t *key, size_t key_size, size_t *key_length )
@@ -180,72 +173,66 @@
     const uint8_t *key, size_t key_length,
     uint8_t *data, size_t data_size, size_t *data_length )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
-    if( psa_key_id_is_builtin(
-            MBEDTLS_SVC_KEY_ID_GET_KEY_ID( psa_get_key_id( attributes ) ) ) )
+    if( key_length == sizeof( psa_drv_slot_number_t ) )
     {
-        if( key_length != sizeof( psa_drv_slot_number_t ) )
-            return( PSA_ERROR_INVALID_ARGUMENT );
+        /* Assume this is a builtin key based on the key material length. */
+        psa_drv_slot_number_t slot_number = *( ( psa_drv_slot_number_t* ) key );
 
-        if( memcmp( key, &ecdsa_slot, sizeof( psa_drv_slot_number_t ) ) == 0 )
+        switch( slot_number )
         {
-            /* This is the ECDSA slot. Verify key attributes before returning
-             * the private key. */
-            if( psa_get_key_type( attributes ) !=
-                PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_bits( attributes ) != 256 )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_algorithm( attributes ) !=
-                PSA_ALG_ECDSA( PSA_ALG_ANY_HASH ) )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( ( psa_get_key_usage_flags( attributes ) &
-                  PSA_KEY_USAGE_EXPORT ) == 0 )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
+            case PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT:
+                /* This is the ECDSA slot. Verify the key's attributes before
+                 * returning the private key. */
+                if( psa_get_key_type( attributes ) !=
+                    PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_bits( attributes ) != 256 )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_algorithm( attributes ) !=
+                    PSA_ALG_ECDSA( PSA_ALG_ANY_HASH ) )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( ( psa_get_key_usage_flags( attributes ) &
+                      PSA_KEY_USAGE_EXPORT ) == 0 )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
 
-            if( data_size < sizeof( test_driver_ecdsa_key ) )
-                return( PSA_ERROR_BUFFER_TOO_SMALL );
+                if( data_size < sizeof( test_driver_ecdsa_key ) )
+                    return( PSA_ERROR_BUFFER_TOO_SMALL );
 
-            memcpy( data, test_driver_ecdsa_key,
-                    sizeof( test_driver_ecdsa_key ) );
-            *data_length = sizeof( test_driver_ecdsa_key );
-            return( PSA_SUCCESS );
+                memcpy( data, test_driver_ecdsa_key,
+                        sizeof( test_driver_ecdsa_key ) );
+                *data_length = sizeof( test_driver_ecdsa_key );
+                return( PSA_SUCCESS );
+
+            case PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT:
+                /* This is the AES slot. Verify the key's attributes before
+                 * returning the key. */
+                if( psa_get_key_type( attributes ) != PSA_KEY_TYPE_AES )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_bits( attributes ) != 128 )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_algorithm( attributes ) != PSA_ALG_CTR )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( ( psa_get_key_usage_flags( attributes ) &
+                      PSA_KEY_USAGE_EXPORT ) == 0 )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+
+                if( data_size < sizeof( test_driver_aes_key ) )
+                    return( PSA_ERROR_BUFFER_TOO_SMALL );
+
+                memcpy( data, test_driver_aes_key,
+                        sizeof( test_driver_aes_key ) );
+                *data_length = sizeof( test_driver_aes_key );
+                return( PSA_SUCCESS );
+
+            default:
+                return( PSA_ERROR_DOES_NOT_EXIST );
         }
-
-        if( memcmp( key, &aes_slot, sizeof( psa_drv_slot_number_t ) ) == 0 )
-        {
-            /* This is the AES slot. Verify key attributes before returning
-             * the key. */
-            if( psa_get_key_type( attributes ) != PSA_KEY_TYPE_AES )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_bits( attributes ) != 128 )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_algorithm( attributes ) != PSA_ALG_CTR )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( ( psa_get_key_usage_flags( attributes ) &
-                  PSA_KEY_USAGE_EXPORT ) == 0 )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-
-            if( data_size < sizeof( test_driver_aes_key ) )
-                return( PSA_ERROR_BUFFER_TOO_SMALL );
-
-            memcpy( data, test_driver_aes_key, sizeof( test_driver_aes_key ) );
-            *data_length = sizeof( test_driver_aes_key );
-            return( PSA_SUCCESS );
-        }
-
-        /* Potentially add more slots here */
-
-        return( PSA_ERROR_DOES_NOT_EXIST );
     }
-#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
-    (void) attributes;
-    (void) key;
-    (void) key_length;
-    (void) data;
-    (void) data_size;
-    (void) data_length;
-    return( PSA_ERROR_NOT_SUPPORTED );
+    else
+    {
+        /* Test driver does not support generic opaque key handling yet. */
+        return( PSA_ERROR_NOT_SUPPORTED );
+    }
 }
 
 psa_status_t test_transparent_export_public_key(
@@ -308,47 +295,41 @@
     const uint8_t *key, size_t key_length,
     uint8_t *data, size_t data_size, size_t *data_length )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
-    if( psa_key_id_is_builtin(
-            MBEDTLS_SVC_KEY_ID_GET_KEY_ID( psa_get_key_id( attributes ) ) ) )
+    if( key_length == sizeof( psa_drv_slot_number_t ) )
     {
-        if( key_length != sizeof( psa_drv_slot_number_t ) )
-            return( PSA_ERROR_INVALID_ARGUMENT );
-
-        if( memcmp( key, &ecdsa_slot, sizeof( psa_drv_slot_number_t ) ) == 0 )
+        /* Assume this is a builtin key based on the key material length. */
+        psa_drv_slot_number_t slot_number = *( ( psa_drv_slot_number_t* ) key );
+        switch( slot_number )
         {
-            /* This is the ECDSA slot. Verify key attributes before returning
-             * the public key. */
-            if( psa_get_key_type( attributes ) !=
-                PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_bits( attributes ) != 256 )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
-            if( psa_get_key_algorithm( attributes ) !=
-                PSA_ALG_ECDSA( PSA_ALG_ANY_HASH ) )
-                return( PSA_ERROR_CORRUPTION_DETECTED );
+            case PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT:
+                /* This is the ECDSA slot. Verify the key's attributes before
+                 * returning the public key. */
+                if( psa_get_key_type( attributes ) !=
+                    PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_bits( attributes ) != 256 )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
+                if( psa_get_key_algorithm( attributes ) !=
+                    PSA_ALG_ECDSA( PSA_ALG_ANY_HASH ) )
+                    return( PSA_ERROR_CORRUPTION_DETECTED );
 
-            if( data_size < sizeof( test_driver_ecdsa_pubkey ) )
-                return( PSA_ERROR_BUFFER_TOO_SMALL );
+                if( data_size < sizeof( test_driver_ecdsa_pubkey ) )
+                    return( PSA_ERROR_BUFFER_TOO_SMALL );
 
-            memcpy( data, test_driver_ecdsa_pubkey,
-                    sizeof( test_driver_ecdsa_pubkey ) );
-            *data_length = sizeof( test_driver_ecdsa_pubkey );
-            return( PSA_SUCCESS );
+                memcpy( data, test_driver_ecdsa_pubkey,
+                        sizeof( test_driver_ecdsa_pubkey ) );
+                *data_length = sizeof( test_driver_ecdsa_pubkey );
+                return( PSA_SUCCESS );
+
+            default:
+                return( PSA_ERROR_DOES_NOT_EXIST );
         }
-
-        /* Potentially add more slots here */
-
-        return( PSA_ERROR_DOES_NOT_EXIST );
     }
-#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
-    (void) attributes;
-    (void) key;
-    (void) key_length;
-    (void) data;
-    (void) data_size;
-    (void) data_length;
-    return( PSA_ERROR_NOT_SUPPORTED );
+    else
+    {
+        /* Test driver does not support generic opaque key handling yet. */
+        return( PSA_ERROR_NOT_SUPPORTED );
+    }
 }
 
 /* The opaque test driver exposes two built-in keys when builtin key support is
@@ -366,13 +347,12 @@
     psa_key_attributes_t *attributes,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+    if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+
     switch( slot_number )
     {
         case PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT:
-            if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
-                return( PSA_ERROR_BUFFER_TOO_SMALL );
-
             psa_set_key_type( attributes, PSA_KEY_TYPE_AES );
             psa_set_key_bits( attributes, 128 );
             psa_set_key_usage_flags(
@@ -387,9 +367,6 @@
             *key_buffer_length = sizeof( psa_drv_slot_number_t );
             return( PSA_SUCCESS );
         case PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT:
-            if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
-                return( PSA_ERROR_BUFFER_TOO_SMALL );
-
             psa_set_key_type(
                 attributes,
                 PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
@@ -407,21 +384,8 @@
             *key_buffer_length = sizeof( psa_drv_slot_number_t );
             return( PSA_SUCCESS );
         default:
-            (void) slot_number;
-            (void) attributes;
-            (void) key_buffer;
-            (void) key_buffer_size;
-            (void) key_buffer_length;
-            return( PSA_ERROR_INVALID_ARGUMENT );
+            return( PSA_ERROR_DOES_NOT_EXIST );
     }
-#else /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
-    (void) slot_number;
-    (void) attributes;
-    (void) key_buffer;
-    (void) key_buffer_size;
-    (void) key_buffer_length;
-    return( PSA_ERROR_DOES_NOT_EXIST );
-#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
 }
 
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */