diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index ab09fa0..cb5d202 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -2994,6 +2994,10 @@
         PSA_ECC_FAMILY_SECP_R1, 256);
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     unsigned char *input_buffer = NULL;
+    const uint8_t server_id[] = PSA_JPAKE_SERVER_ID;
+    const uint8_t client_id[] = PSA_JPAKE_CLIENT_ID;
+    const size_t server_id_len = strlen(PSA_JPAKE_SERVER_ID);
+    const size_t client_id_len = strlen(PSA_JPAKE_CLIENT_ID);
     const size_t size_key_share = PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                                       PSA_PAKE_STEP_KEY_SHARE);
     unsigned char *output_buffer = NULL;
@@ -3035,8 +3039,8 @@
     TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
                PSA_SUCCESS);
 
-    TEST_EQUAL(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER),
-               PSA_SUCCESS);
+    PSA_ASSERT(psa_pake_set_user(&operation, server_id, server_id_len));
+    PSA_ASSERT(psa_pake_set_peer(&operation, client_id, client_id_len));
 
     TEST_EQUAL(psa_pake_set_password_key(&operation, key),
                PSA_SUCCESS);
@@ -3184,6 +3188,10 @@
         PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_derivation_operation_t client_derive =
         PSA_KEY_DERIVATION_OPERATION_INIT;
+    const uint8_t server_id[] = PSA_JPAKE_SERVER_ID;
+    const uint8_t client_id[] = PSA_JPAKE_CLIENT_ID;
+    const size_t server_id_len = strlen(PSA_JPAKE_SERVER_ID);
+    const size_t client_id_len = strlen(PSA_JPAKE_CLIENT_ID);
     pake_in_driver = in_driver;
     /* driver setup is called indirectly through pake_output/pake_input */
     if (pake_in_driver) {
@@ -3229,9 +3237,11 @@
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
 
 
-    PSA_ASSERT(psa_pake_set_role(&server, PSA_PAKE_ROLE_SERVER));
+    PSA_ASSERT(psa_pake_set_user(&server, server_id, server_id_len));
+    PSA_ASSERT(psa_pake_set_peer(&server, client_id, client_id_len));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
-    PSA_ASSERT(psa_pake_set_role(&client, PSA_PAKE_ROLE_CLIENT));
+    PSA_ASSERT(psa_pake_set_user(&client, client_id, client_id_len));
+    PSA_ASSERT(psa_pake_set_peer(&client, server_id, server_id_len));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);