diff --git a/tests/include/drivers/keygen.h b/tests/include/drivers/keygen.h
index 436df34..7e6ed27 100644
--- a/tests/include/drivers/keygen.h
+++ b/tests/include/drivers/keygen.h
@@ -28,7 +28,7 @@
 #include MBEDTLS_CONFIG_FILE
 #endif
 
-#if defined(MBEDTLS_TEST_HOOKS)
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVER_TEST)
 #include <psa/crypto_driver_common.h>
 
 extern void *test_driver_keygen_forced_output;
@@ -45,5 +45,5 @@
     const psa_key_attributes_t *attributes,
     uint8_t *key, size_t key_size, size_t *key_length );
 
-#endif /* MBEDTLS_TEST_HOOKS */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVER_TEST */
 #endif /* MBEDTLS_PSA_CRYPTO_TEST_DRIVERS_KEYGEN_H */