diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 779f594..60befa7 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -31,6 +31,29 @@
 #define ASSERT_OPERATION_IS_ACTIVE(   operation ) TEST_ASSERT( operation.id != 0 )
 #define ASSERT_OPERATION_IS_INACTIVE( operation ) TEST_ASSERT( operation.id == 0 )
 
+#if defined(PSA_WANT_ALG_JPAKE)
+void ecjpake_operation_setup( psa_pake_operation_t *operation,
+                              psa_pake_cipher_suite_t *cipher_suite,
+                              psa_pake_role_t role,
+                              mbedtls_svc_key_id_t key,
+                              size_t key_available )
+{
+    *operation = psa_pake_operation_init();
+
+    TEST_EQUAL( psa_pake_setup( operation, cipher_suite ),
+                PSA_SUCCESS );
+
+    TEST_EQUAL( psa_pake_set_role( operation, role),
+                PSA_SUCCESS );
+
+    if( key_available )
+        TEST_EQUAL( psa_pake_set_password_key( operation, key ),
+                    PSA_SUCCESS );
+exit:
+    return;
+}
+#endif
+
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
 
@@ -8740,7 +8763,6 @@
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
-    psa_pake_operation_t op_copy = psa_pake_operation_init();
     psa_algorithm_t alg = alg_arg;
     psa_pake_primitive_t primitive = primitive_arg;
     psa_key_type_t key_type_pw = key_type_pw_arg;
@@ -8839,22 +8861,25 @@
     if( input_first )
     {
         /* Invalid parameters (input) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                     NULL, 0 ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid parameters (step) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_BAD_STATE );
 
         /* Possibly valid */
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len );
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                     output_buffer, size_key_share ),
                     expected_status_input_output);
@@ -8875,22 +8900,25 @@
     else
     {
         /* Invalid parameters (output) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                      NULL, 0, NULL ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        op_copy = operation;
         /* Invalid parameters (step) */
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                key, pw_data->len );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_BAD_STATE );
 
         /* Possibly valid */
+        ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len );
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                      output_buffer, buf_size, &output_len ),
                     expected_status_input_output );