diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 1b144df..49887b5 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8740,6 +8740,7 @@
     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_pake_primitive_t primitive = primitive_arg;
     psa_key_type_t key_type_pw = key_type_pw_arg;
     psa_key_usage_t key_usage_pw = key_usage_pw_arg;
     psa_algorithm_t hash_alg = hash_arg;
@@ -8757,9 +8758,9 @@
 
     PSA_INIT( );
 
-    ASSERT_ALLOC( output_buffer,
-                  PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg,
-                                       PSA_PAKE_STEP_KEY_SHARE) );
+    size_t buf_size = PSA_PAKE_OUTPUT_SIZE(alg, primitive_arg,
+                                       PSA_PAKE_STEP_KEY_SHARE);
+    ASSERT_ALLOC( output_buffer, buf_size );
 
     if( pw_data->len > 0 )
     {
@@ -8771,7 +8772,7 @@
     }
 
     psa_pake_cs_set_algorithm( &cipher_suite, alg );
-    psa_pake_cs_set_primitive( &cipher_suite, primitive_arg );
+    psa_pake_cs_set_primitive( &cipher_suite, primitive );
     psa_pake_cs_set_hash( &cipher_suite, hash_alg );
 
     PSA_ASSERT( psa_pake_abort( &operation ) );
@@ -8825,6 +8826,13 @@
     TEST_EQUAL( psa_pake_set_peer( &operation, unsupported_id, 4 ),
                 PSA_ERROR_NOT_SUPPORTED );
 
+    const size_t size_key_share = PSA_PAKE_INPUT_SIZE( alg, primitive,
+                                                PSA_PAKE_STEP_KEY_SHARE );
+    const size_t size_zk_public = PSA_PAKE_INPUT_SIZE( alg, primitive,
+                                                PSA_PAKE_STEP_ZK_PUBLIC );
+    const size_t size_zk_proof = PSA_PAKE_INPUT_SIZE( alg, primitive,
+                                                PSA_PAKE_STEP_ZK_PROOF );
+
     /* First round */
     if( input_first )
     {
@@ -8833,28 +8841,23 @@
                                     NULL, 0 ),
                     PSA_ERROR_INVALID_ARGUMENT );
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
-                                    output_buffer, 66 ),
+                                    output_buffer, size_zk_proof ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
-                                    output_buffer, 66 ),
+                                    output_buffer, size_zk_proof ),
                     PSA_ERROR_BAD_STATE );
 
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
-                                    output_buffer, 66 ),
+                                    output_buffer, size_key_share ),
                     expected_status_input_output);
 
         if( expected_status_input_output == 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 );
+                                    output_buffer, size_zk_public + 1 ),
+                        PSA_ERROR_INVALID_ARGUMENT );
         }
     }
     else
@@ -8864,15 +8867,15 @@
                                      NULL, 0, NULL ),
                     PSA_ERROR_INVALID_ARGUMENT );
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
-                                     output_buffer, 512, &output_len ),
+                                     output_buffer, buf_size, &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 ),
+                                     output_buffer, buf_size, &output_len ),
                     PSA_ERROR_BAD_STATE );
 
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
-                                     output_buffer, 512, &output_len ),
+                                     output_buffer, buf_size, &output_len ),
                     expected_status_input_output );
 
         if( expected_status_input_output == PSA_SUCCESS )
@@ -8881,13 +8884,8 @@
 
             /* Buffer too small */
             TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
-                                         output_buffer, 5, &output_len ),
+                                         output_buffer, size_zk_public - 1, &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 );
         }
     }