diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 799ec04d3..34c416326 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6448,47 +6448,47 @@ persistent_key_load_key_from_storage:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b
 
 PSA PAKE: invalid alg
 depends_on: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
+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:0:"abcd":PSA_ERROR_INVALID_ARGUMENT
 
 PSA PAKE: invalid primitive type
 depends_on: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
+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:0:"abcd":PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive family
 depends_on: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
+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:0:"abcd":PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive bits
 depends_on: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
+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:0:"abcd":PSA_ERROR_NOT_SUPPORTED
 
-PSA PAKE: ecjpake setup server
+PSA PAKE: ecjpake setup server output step first
 depends_on: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
+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:0:"abcd":0
+
+PSA PAKE: ecjpake setup server input step first
+depends_on: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:1:"abcd":0
 
 PSA PAKE: ecjpake setup server empty password
 depends_on: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
+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:0:"":PSA_ERROR_BAD_STATE
 
-PSA PAKE: ecjpake setup server invalid step
+PSA PAKE: ecjpake setup client output step first
 depends_on: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
+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:0:"abcd":0
 
-PSA PAKE: ecjpake setup client
+PSA PAKE: ecjpake setup client input step first
 depends_on: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
+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:1:"abcd":0
 
 PSA PAKE: ecjpake setup client empty password
 depends_on: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_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
+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:0:"":PSA_ERROR_BAD_STATE
 
 PSA PAKE: ecjpake setup invalid role NONE
 depends_on: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
+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:0:"abcd":PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: ecjpake rounds
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 66572670a..e091ad405 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8094,7 +8094,7 @@ exit:
 
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_ECJPAKE */
 void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
-                    int output_step_arg, data_t *pw_data,
+                    int input_first, data_t *pw_data,
                     int expected_status_arg )
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
@@ -8102,7 +8102,6 @@ void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
     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;
@@ -8113,7 +8112,8 @@ void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
     PSA_INIT( );
 
     ASSERT_ALLOC( output_buffer,
-                  PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg, step) );
+                  PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg,
+                                       PSA_PAKE_STEP_KEY_SHARE) );
 
     if( pw_data->len > 0 )
     {
@@ -8138,9 +8138,10 @@ void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
                 PSA_ERROR_BAD_STATE );
     TEST_EQUAL( psa_pake_set_role( &operation, role ),
                 PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_output( &operation, step, NULL, 0, NULL ),
+    TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                                 NULL, 0, NULL ),
                 PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_input( &operation, step, NULL, 0),
+    TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE, NULL, 0),
                 PSA_ERROR_BAD_STATE );
 
     PSA_ASSERT( psa_pake_abort( &operation ) );
@@ -8190,21 +8191,71 @@ void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
     TEST_EQUAL( psa_pake_set_peer( &operation, unsupported_id, 4 ),
                 PSA_ERROR_NOT_SUPPORTED );
 
-    /* First round Output */
-    status = psa_pake_output( &operation, step, output_buffer,
-                              512, &output_len );
-    if( status != PSA_SUCCESS )
+    /* First round */
+    if( input_first )
     {
-        TEST_EQUAL( status, expected_status );
-        goto exit;
+        /* Invalid parameters */
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                    NULL, 0 ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+                                    output_buffer, 66 ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        /* Invalid first step */
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                    output_buffer, 66 ),
+                    PSA_ERROR_BAD_STATE );
+
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                                    output_buffer, 66 ),
+                    expected_status);
+
+        if( expected_status == PSA_SUCCESS )
+        {
+            /* Buffer too large */
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                    output_buffer, 512 ),
+                        PSA_ERROR_INSUFFICIENT_MEMORY );
+
+            /* The operation should be aborted at this point */
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, 66 ),
+                        PSA_ERROR_BAD_STATE );
+        }
     }
     else
-        PSA_ASSERT( status );
+    {
+        /* Invalid parameters */
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                     NULL, 0, NULL ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+                                     output_buffer, 512, &output_len ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        /* Invalid first step */
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                     output_buffer, 512, &output_len ),
+                    PSA_ERROR_BAD_STATE );
 
-    TEST_EQUAL( psa_pake_output( &operation, step, NULL, 0, NULL),
-                PSA_ERROR_INVALID_ARGUMENT );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                                     output_buffer, 512, &output_len ),
+                    expected_status );
 
-    TEST_ASSERT( output_len > 0 );
+        if( expected_status == PSA_SUCCESS )
+        {
+            TEST_ASSERT( output_len > 0 );
+
+            /* Buffer too small */
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                         output_buffer, 5, &output_len ),
+                        PSA_ERROR_BUFFER_TOO_SMALL );
+
+            /* The operation should be aborted at this point */
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                         output_buffer, 512, &output_len ),
+                        PSA_ERROR_BAD_STATE );
+        }
+    }
 
 exit:
     PSA_ASSERT( psa_destroy_key( key ) );