diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index 5ee401b..b1d3d44 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -50,27 +50,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_setup(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite);
-
-psa_status_t mbedtls_test_transparent_set_password_key(
-    const psa_key_attributes_t *attributes,
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    uint8_t *key_buffer,
-    size_t key_size);
-
-psa_status_t mbedtls_test_transparent_pake_set_user(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *user_id,
-    size_t user_id_len);
-
-psa_status_t mbedtls_test_transparent_pake_set_peer(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *peer_id,
-    size_t peer_id_len);
-
-psa_status_t mbedtls_test_transparent_pake_set_role(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_pake_role_t role);
+    const psa_crypto_driver_pake_inputs_t *inputs);
 
 psa_status_t mbedtls_test_transparent_pake_output(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
@@ -94,7 +74,7 @@
 
 psa_status_t mbedtls_test_opaque_pake_setup(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite);
+    const psa_crypto_driver_pake_inputs_t *inputs);
 
 psa_status_t mbedtls_test_opaque_set_password_key(
     const psa_key_attributes_t *attributes,