diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index 81e8711..5ee401b 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -87,7 +87,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_get_implicit_key(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_key_derivation_operation_t *output);
+    uint8_t *output, size_t *output_size);
 
 psa_status_t mbedtls_test_transparent_pake_abort(
     mbedtls_transparent_test_driver_pake_operation_t *operation);
@@ -131,7 +131,7 @@
 
 psa_status_t mbedtls_test_opaque_pake_get_implicit_key(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    psa_key_derivation_operation_t *output);
+    uint8_t *output, size_t *output_size);
 
 psa_status_t mbedtls_test_opaque_pake_abort(
     mbedtls_opaque_test_driver_pake_operation_t *operation);