diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 262070c6a..f905ea4f9 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -273,7 +273,7 @@ static psa_status_t get_expected_key_size( const psa_key_attributes_t *attribute
 
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
-#ifdef TEST_KEY_CONTEXT_SIZE_FUNCTION
+#ifdef TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
             *expected_size = test_size_function( key_type, key_bits );
             return( PSA_SUCCESS );
 #else /* TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */