diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 1182c00..298a5af 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6445,3 +6445,51 @@
 
 PSA derive persistent key: HKDF SHA-256, exportable
 persistent_key_load_key_from_storage:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_TYPE_RAW_DATA:1024:PSA_KEY_USAGE_EXPORT:0:DERIVE_KEY
+
+PSA PAKE: invalid alg
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_SHA_256:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"abcd":PSA_ERROR_INVALID_ARGUMENT
+
+PSA PAKE: invalid primitive type
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_DH, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"abcd":PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: invalid primitive family
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_K1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"abcd":PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: invalid primitive bits
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 128):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"abcd":PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: ecjpake setup server
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"abcd":0
+
+PSA PAKE: ecjpake setup server empty password
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_KEY_SHARE:"":PSA_ERROR_BAD_STATE
+
+PSA PAKE: ecjpake setup server invalid step
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:PSA_PAKE_STEP_ZK_PROOF:"abcd":PSA_ERROR_BAD_STATE
+
+PSA PAKE: ecjpake setup client
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:PSA_PAKE_STEP_KEY_SHARE:"abcd":0
+
+PSA PAKE: ecjpake setup client empty password
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:PSA_PAKE_STEP_KEY_SHARE:"":PSA_ERROR_BAD_STATE
+
+PSA PAKE: ecjpake setup client invalid step
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:PSA_PAKE_STEP_ZK_PROOF:"abcd":PSA_ERROR_BAD_STATE
+
+PSA PAKE: ecjpake setup invalid role NONE
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:PSA_PAKE_STEP_KEY_SHARE:"abcd":PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: ecjpake rounds
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef"
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 7d368cf..9b7bb20 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8091,3 +8091,327 @@
     PSA_DONE();
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_ECJPAKE_C */
+void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
+                    int output_step_arg, data_t *pw_data,
+                    int expected_status_arg )
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_algorithm_t alg = alg_arg;
+    psa_algorithm_t hash_alg = hash_arg;
+    psa_pake_role_t role = role_arg;
+    psa_pake_step_t step = output_step_arg;
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t expected_status = expected_status_arg;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    unsigned char *output_buffer = NULL;
+    size_t output_len = 0;
+
+    PSA_INIT( );
+
+    ASSERT_ALLOC( output_buffer,
+                  PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg, step) );
+
+    if( pw_data->len > 0 )
+    {
+        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+        psa_set_key_algorithm( &attributes, alg );
+        psa_set_key_type( &attributes, PSA_KEY_TYPE_PASSWORD );
+        PSA_ASSERT( psa_import_key( &attributes, pw_data->x, pw_data->len,
+                                    &key ) );
+    }
+
+    psa_pake_cs_set_algorithm( &cipher_suite, alg );
+    psa_pake_cs_set_primitive( &cipher_suite, primitive_arg );
+    psa_pake_cs_set_hash( &cipher_suite, hash_alg );
+
+    status = psa_pake_setup( &operation, &cipher_suite );
+    if( status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( status, expected_status );
+        goto exit;
+    }
+    else
+        PSA_ASSERT( status );
+
+    status = psa_pake_set_role( &operation, role );
+    if( status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( status, expected_status );
+        goto exit;
+    }
+    else
+        PSA_ASSERT( status );
+
+    if( pw_data->len > 0 )
+    {
+        status = psa_pake_set_password_key( &operation, key );
+        if( status != PSA_SUCCESS )
+        {
+            TEST_EQUAL( status, expected_status );
+            goto exit;
+        }
+        else
+            PSA_ASSERT( status );
+    }
+
+    /* First round Output */
+    status = psa_pake_output( &operation, step, output_buffer,
+                              512, &output_len );
+    if( status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( status, expected_status );
+        goto exit;
+    }
+    else
+        PSA_ASSERT( status );
+
+    TEST_ASSERT( output_len > 0 );
+
+exit:
+    PSA_ASSERT( psa_destroy_key( key ) );
+    PSA_ASSERT( psa_pake_abort( &operation ) );
+    mbedtls_free( output_buffer );
+    PSA_DONE( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_ECJPAKE_C */
+void ecjpake_rounds( int alg_arg, int primitive_arg, int hash_arg,
+                     int derive_alg_arg, data_t *pw_data )
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t server = psa_pake_operation_init();
+    psa_pake_operation_t client = psa_pake_operation_init();
+    psa_algorithm_t alg = alg_arg;
+    psa_algorithm_t hash_alg = hash_arg;
+    psa_algorithm_t derive_alg = derive_alg_arg;
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_derivation_operation_t server_derive =
+                            PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_key_derivation_operation_t client_derive =
+                            PSA_KEY_DERIVATION_OPERATION_INIT;
+    unsigned char *buffer0 = NULL, *buffer1 = NULL;
+    size_t buffer_length = (
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg, PSA_PAKE_STEP_KEY_SHARE) +
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg, PSA_PAKE_STEP_ZK_PUBLIC) +
+        PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg, PSA_PAKE_STEP_ZK_PROOF)) * 2;
+    size_t buffer0_off = 0;
+    size_t buffer1_off = 0;
+    size_t s_g1_len, s_g2_len, s_a_len;
+    size_t s_g1_off, s_g2_off, s_a_off;
+    size_t s_x1_pk_len, s_x2_pk_len, s_x2s_pk_len;
+    size_t s_x1_pk_off, s_x2_pk_off, s_x2s_pk_off;
+    size_t s_x1_pr_len, s_x2_pr_len, s_x2s_pr_len;
+    size_t s_x1_pr_off, s_x2_pr_off, s_x2s_pr_off;
+    size_t c_g1_len, c_g2_len, c_a_len;
+    size_t c_g1_off, c_g2_off, c_a_off;
+    size_t c_x1_pk_len, c_x2_pk_len, c_x2s_pk_len;
+    size_t c_x1_pk_off, c_x2_pk_off, c_x2s_pk_off;
+    size_t c_x1_pr_len, c_x2_pr_len, c_x2s_pr_len;
+    size_t c_x1_pr_off, c_x2_pr_off, c_x2s_pr_off;
+
+    PSA_INIT( );
+
+    ASSERT_ALLOC( buffer0, buffer_length );
+    ASSERT_ALLOC( buffer1, buffer_length );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_PASSWORD );
+    PSA_ASSERT( psa_import_key( &attributes, pw_data->x, pw_data->len,
+                &key ) );
+
+    psa_pake_cs_set_algorithm( &cipher_suite, alg );
+    psa_pake_cs_set_primitive( &cipher_suite, primitive_arg );
+    psa_pake_cs_set_hash( &cipher_suite, hash_alg );
+
+    PSA_ASSERT( psa_pake_setup( &server, &cipher_suite ) );
+    PSA_ASSERT( psa_pake_setup( &client, &cipher_suite ) );
+
+    PSA_ASSERT( psa_pake_set_role( &server, PSA_PAKE_ROLE_SERVER ) );
+    PSA_ASSERT( psa_pake_set_role( &client, PSA_PAKE_ROLE_CLIENT ) );
+
+    PSA_ASSERT( psa_pake_set_password_key( &server, key ) );
+    PSA_ASSERT( psa_pake_set_password_key( &client, key ) );
+
+    /* Server first round Output */
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_g1_len ) );
+    s_g1_off = buffer0_off;
+    buffer0_off += s_g1_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x1_pk_len ) );
+    s_x1_pk_off = buffer0_off;
+    buffer0_off += s_x1_pk_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x1_pr_len ) );
+    s_x1_pr_off = buffer0_off;
+    buffer0_off += s_x1_pr_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_g2_len ) );
+    s_g2_off = buffer0_off;
+    buffer0_off += s_g2_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x2_pk_len ) );
+    s_x2_pk_off = buffer0_off;
+    buffer0_off += s_x2_pk_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x2_pr_len ) );
+    s_x2_pr_off = buffer0_off;
+    buffer0_off += s_x2_pr_len;
+
+    /* Client first round Output */
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_g1_len ) );
+    c_g1_off = buffer1_off;
+    buffer1_off += c_g1_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x1_pk_len ) );
+    c_x1_pk_off = buffer1_off;
+    buffer1_off += c_x1_pk_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x1_pr_len ) );
+    c_x1_pr_off = buffer1_off;
+    buffer1_off += c_x1_pr_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_g2_len ) );
+    c_g2_off = buffer1_off;
+    buffer1_off += c_g2_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x2_pk_len ) );
+    c_x2_pk_off = buffer1_off;
+    buffer1_off += c_x2_pk_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x2_pr_len ) );
+    c_x2_pr_off = buffer1_off;
+    buffer1_off += c_x2_pr_len;
+
+    /* Client first round Input */
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer0 + s_g1_off, s_g1_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer0 + s_x1_pk_off, s_x1_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer0 + s_x1_pr_off, s_x1_pr_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer0 + s_g2_off, s_g2_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer0 + s_x2_pk_off, s_x2_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer0 + s_x2_pr_off, s_x2_pr_len ) );
+
+    /* Server first round Input */
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer1 + c_g1_off, c_g1_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer1 + c_x1_pk_off, c_x1_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer1 + c_x1_pr_off, c_x1_pr_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer1 + c_g2_off, c_g2_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer1 + c_x2_pk_off, c_x2_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer1 + c_x2_pr_off, c_x2_pr_len ) );
+
+    /* Server second round Output */
+    buffer0_off = 0;
+
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_a_len ) );
+    s_a_off = buffer0_off;
+    buffer0_off += s_a_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x2s_pk_len ) );
+    s_x2s_pk_off = buffer0_off;
+    buffer0_off += s_x2s_pk_len;
+    PSA_ASSERT( psa_pake_output( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer0 + buffer0_off,
+                                 512 - buffer0_off, &s_x2s_pr_len ) );
+    s_x2s_pr_off = buffer0_off;
+    buffer0_off += s_x2s_pr_len;
+
+    /* Client second round Output */
+    buffer1_off = 0;
+
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_a_len ) );
+    c_a_off = buffer1_off;
+    buffer1_off += c_a_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x2s_pk_len ) );
+    c_x2s_pk_off = buffer1_off;
+    buffer1_off += c_x2s_pk_len;
+    PSA_ASSERT( psa_pake_output( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                 buffer1 + buffer1_off,
+                                 512 - buffer1_off, &c_x2s_pr_len ) );
+    c_x2s_pr_off = buffer1_off;
+    buffer1_off += c_x2s_pr_len;
+
+    /* Client second round Input */
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer0 + s_a_off, s_a_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer0 + s_x2s_pk_off, s_x2s_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &client, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer0 + s_x2s_pr_off, s_x2s_pr_len ) );
+
+    /* Server second round Input */
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer1 + c_a_off, c_a_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PUBLIC,
+                                buffer1 + c_x2s_pk_off, c_x2s_pk_len ) );
+    PSA_ASSERT( psa_pake_input( &server, PSA_PAKE_STEP_ZK_PROOF,
+                                buffer1 + c_x2s_pr_off, c_x2s_pr_len ) );
+
+
+    /* Get shared key */
+    PSA_ASSERT( psa_key_derivation_setup( &server_derive, derive_alg ) );
+    PSA_ASSERT( psa_key_derivation_setup( &client_derive, derive_alg ) );
+
+    if( PSA_ALG_IS_TLS12_PRF( derive_alg ) ||
+        PSA_ALG_IS_TLS12_PSK_TO_MS( derive_alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( &server_derive,
+                                                PSA_KEY_DERIVATION_INPUT_SEED,
+                                                (const uint8_t*) "", 0) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &client_derive,
+                                                PSA_KEY_DERIVATION_INPUT_SEED,
+                                                (const uint8_t*) "", 0) );
+    }
+
+    PSA_ASSERT( psa_pake_get_implicit_key( &server, &server_derive ) );
+    PSA_ASSERT( psa_pake_get_implicit_key( &client, &client_derive ) );
+
+exit:
+    psa_key_derivation_abort( &server_derive );
+    psa_key_derivation_abort( &client_derive );
+    psa_destroy_key( key );
+    psa_pake_abort( &server );
+    psa_pake_abort( &client );
+    mbedtls_free( buffer0 );
+    mbedtls_free( buffer1 );
+    PSA_DONE( );
+}
+/* END_CASE */