test_driver_extension: manage public and private keys the same way
Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
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