diff --git a/tests/src/drivers/keygen.c b/tests/src/drivers/keygen.c
index c883e00..1f96fc8 100644
--- a/tests/src/drivers/keygen.c
+++ b/tests/src/drivers/keygen.c
@@ -25,7 +25,7 @@
 #include MBEDTLS_CONFIG_FILE
 #endif
 
-#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(MBEDTLS_PSA_CRYPTO_DRIVER_TEST)
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
 #include "psa/crypto.h"
 #include "mbedtls/ecp.h"
 #include "mbedtls/error.h"
@@ -126,4 +126,4 @@
     return( PSA_ERROR_NOT_SUPPORTED );
 }
 
-#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && MBEDTLS_PSA_CRYPTO_DRIVER_TEST */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */