crypto_config_test_driver_extension: treat PUBLIC_KEY the same way as KEY_PAIR

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 91ee320..a285bc7 100644
--- a/tests/include/test/drivers/crypto_config_test_driver_extension.h
+++ b/tests/include/test/drivers/crypto_config_test_driver_extension.h
@@ -198,6 +198,14 @@
 #endif
 #endif
 
+#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_ECC_PUBLIC_KEY 1
+#endif
+#endif
+
 #if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR)
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR)
 #undef MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR
@@ -206,6 +214,14 @@
 #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_CHACHA20)
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_CHACHA20)
 #undef MBEDTLS_PSA_ACCEL_KEY_TYPE_CHACHA20
@@ -248,6 +264,4 @@
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_HMAC 1
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_ARC4 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