diff --git a/tests/src/drivers/key_management.c b/tests/src/drivers/key_management.c
index 34eb614..a788934 100644
--- a/tests/src/drivers/key_management.c
+++ b/tests/src/drivers/key_management.c
@@ -43,7 +43,8 @@
     const psa_key_attributes_t *attributes,
     uint8_t *key, size_t key_size, size_t *key_length )
 {
-#if !defined(MBEDTLS_PSA_BUILTIN_ECC_KEY_PAIR) && !defined(MBEDTLS_PSA_BUILTIN_ECC_PUBLIC_KEY)
+#if !defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) && \
+    !defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
     (void)attributes;
 #endif /* !MBEDTLS_PSA_BUILTIN_ECC_KEY_PAIR && !MBEDTLS_PSA_BUILTIN_ECC_PUBLIC_KEY */
     ++test_driver_key_management_hooks.hits;
@@ -62,7 +63,8 @@
     }
 
     /* Copied from psa_crypto.c */
-#if defined(MBEDTLS_PSA_BUILTIN_ECC_KEY_PAIR) || defined(MBEDTLS_PSA_BUILTIN_ECC_PUBLIC_KEY)
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
     if ( PSA_KEY_TYPE_IS_ECC( psa_get_key_type( attributes ) )
          && PSA_KEY_TYPE_IS_KEY_PAIR( psa_get_key_type( attributes ) ) )
     {
@@ -118,7 +120,7 @@
         return( status );
     }
     else
-#endif /* MBEDTLS_PSA_BUILTIN_ECC_KEY_PAIR || MBEDTLS_PSA_BUILTIN_ECC_PUBLIC_KEY */
+#endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR || MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY */
     return( PSA_ERROR_NOT_SUPPORTED );
 }
 
@@ -143,7 +145,8 @@
     if( test_driver_key_management_hooks.forced_status != PSA_SUCCESS )
         return( test_driver_key_management_hooks.forced_status );
 
-#if defined(MBEDTLS_ECP_C)
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
     psa_key_type_t type = psa_get_key_type( attributes );
     if ( PSA_KEY_TYPE_IS_ECC( type ) )
     {
@@ -234,11 +237,12 @@
     }
     return( PSA_ERROR_NOT_SUPPORTED );
 #else
+    (void) attributes;
     (void) data;
     (void) data_length;
     (void) bits;
     return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* MBEDTLS_ECP_C */
+#endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR || MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY */
 }
 
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */