diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index 1f53008..23cb98a 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -57,16 +57,14 @@
 
 psa_status_t mbedtls_test_transparent_pake_output(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_pake_step_t step,
-    const psa_pake_computation_stage_t *computation_stage,
+    psa_pake_driver_step_t step,
     uint8_t *output,
     size_t output_size,
     size_t *output_length);
 
 psa_status_t mbedtls_test_transparent_pake_input(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_pake_step_t step,
-    const psa_pake_computation_stage_t *computation_stage,
+    psa_pake_driver_step_t step,
     const uint8_t *input,
     size_t input_length);
 
@@ -103,16 +101,14 @@
 
 psa_status_t mbedtls_test_opaque_pake_output(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    psa_pake_step_t step,
-    const psa_pake_computation_stage_t *computation_stage,
+    psa_pake_driver_step_t step,
     uint8_t *output,
     size_t output_size,
     size_t *output_length);
 
 psa_status_t mbedtls_test_opaque_pake_input(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    psa_pake_step_t step,
-    const psa_pake_computation_stage_t *computation_stage,
+    psa_pake_driver_step_t step,
     const uint8_t *input,
     size_t input_length);