diff --git a/tests/include/test/drivers/crypto_config_test_driver_extension.h b/tests/include/test/drivers/crypto_config_test_driver_extension.h
index 4f3e5ad..7555870 100644
--- a/tests/include/test/drivers/crypto_config_test_driver_extension.h
+++ b/tests/include/test/drivers/crypto_config_test_driver_extension.h
@@ -304,11 +304,27 @@
 #endif
 #endif
 
-#if defined(MBEDTLS_PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_LEGACY)
-#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR)
-#undef MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
+#undef MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY
 #else
-#define MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR 1
+#define MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY 1
+#endif
+#endif
+
+#if defined(PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY)
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+#undef MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY
+#else
+#define MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY 1
+#endif
+#endif
+
+#if defined(PSA_WANT_KEY_TYPE_DH_PUBLIC_KEY)
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_DH_PUBLIC_KEY)
+#undef MBEDTLS_PSA_ACCEL_KEY_TYPE_DH_PUBLIC_KEY
+#else
+#define MBEDTLS_PSA_ACCEL_KEY_TYPE_DH_PUBLIC_KEY 1
 #endif
 #endif
 
@@ -379,7 +395,4 @@
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_DERIVE 1
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_HMAC 1
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_DES 1
-#define MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY 1
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_RAW_DATA 1
-#define MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY 1
-#define MBEDTLS_PSA_ACCEL_KEY_TYPE_DH_PUBLIC_KEY 1