diff --git a/tests/src/drivers/keygen.c b/tests/src/drivers/keygen.c
index 4c830c3..c883e00 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_TEST_HOOKS)
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(MBEDTLS_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_TEST_HOOKS */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && MBEDTLS_PSA_CRYPTO_DRIVER_TEST */
diff --git a/tests/src/drivers/signature.c b/tests/src/drivers/signature.c
index 32ebbfd..114007a 100644
--- a/tests/src/drivers/signature.c
+++ b/tests/src/drivers/signature.c
@@ -25,7 +25,7 @@
 #include MBEDTLS_CONFIG_FILE
 #endif
 
-#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(MBEDTLS_TEST_HOOKS)
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(MBEDTLS_PSA_CRYPTO_DRIVER_TEST)
 #include "psa/crypto.h"
 #include "mbedtls/ecp.h"
 
@@ -298,4 +298,4 @@
     return( PSA_ERROR_NOT_SUPPORTED );
 }
 
-#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && MBEDTLS_TEST_HOOKS */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && MBEDTLS_PSA_CRYPTO_DRIVER_TEST */