diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 90948d7..9efee51 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -517,57 +517,76 @@
     return( 0 );
 }
 
+static int setup_key_derivation_wrap( psa_key_derivation_operation_t* operation,
+                                      psa_key_handle_t handle,
+                                      psa_algorithm_t alg,
+                                      unsigned char* input1, size_t input1_length,
+                                      unsigned char* input2, size_t input2_length,
+                                      size_t capacity )
+{
+    PSA_ASSERT( psa_key_derivation_setup( operation, alg ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SALT,
+                                                    input1, input1_length ) );
+        PSA_ASSERT( psa_key_derivation_input_key( operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_INFO,
+                                                    input2,
+                                                    input2_length ) );
+    }
+    else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SEED,
+                                                    input1, input1_length ) );
+        PSA_ASSERT( psa_key_derivation_input_key( operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_LABEL,
+                                                    input2, input2_length ) );
+    }
+    else
+    {
+        TEST_ASSERT( ! "Key derivation algorithm not supported" );
+    }
+
+    PSA_ASSERT( psa_key_derivation_set_capacity( operation, capacity ) );
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+
 static int exercise_key_derivation_key( psa_key_handle_t handle,
                                         psa_key_usage_t usage,
                                         psa_algorithm_t alg )
 {
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    unsigned char label[16] = "This is a label.";
-    size_t label_length = sizeof( label );
-    unsigned char seed[16] = "abcdefghijklmnop";
-    size_t seed_length = sizeof( seed );
+    unsigned char input1[] = "Input 1";
+    size_t input1_length = sizeof( input1 );
+    unsigned char input2[] = "Input 2";
+    size_t input2_length = sizeof( input2 );
     unsigned char output[1];
+    size_t capacity = sizeof( output );
 
     if( usage & PSA_KEY_USAGE_DERIVE )
     {
-        PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-        if( PSA_ALG_IS_HKDF( alg ) )
-        {
-            PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                        PSA_KEY_DERIVATION_INPUT_SALT,
-                                                        label,
-                                                        label_length ) );
-            PSA_ASSERT( psa_key_derivation_input_key( &operation,
-                                                      PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                      handle ) );
-            PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                        PSA_KEY_DERIVATION_INPUT_INFO,
-                                                        seed,
-                                                        seed_length ) );
-        }
-        else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
-                 PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
-        {
-            PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                        PSA_KEY_DERIVATION_INPUT_SEED,
-                                                        seed,
-                                                        seed_length ) );
-            PSA_ASSERT( psa_key_derivation_input_key( &operation,
-                                                      PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                      handle ) );
-            PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                        PSA_KEY_DERIVATION_INPUT_LABEL,
-                                                        label,
-                                                        label_length ) );
-        }
-        else
-        {
-            TEST_ASSERT( ! "Key derivation algorithm not supported" );
-        }
+        if( !setup_key_derivation_wrap( &operation, handle, alg,
+                                        input1, input1_length,
+                                        input2, input2_length, capacity ) )
+            goto exit;
 
         PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
                                                      output,
-                                                     sizeof( output ) ) );
+                                                     capacity ) );
         PSA_ASSERT( psa_key_derivation_abort( &operation ) );
     }
 
@@ -4362,40 +4381,11 @@
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &handle ) );
 
-    PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-    PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
-                                                 requested_capacity ) );
-
-    /* Extraction phase. */
-    if( PSA_ALG_IS_HKDF( alg ) )
-    {
-        PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                    PSA_KEY_DERIVATION_INPUT_SALT,
-                                                    input1->x, input1->len ) );
-        PSA_ASSERT( psa_key_derivation_input_key( &operation,
-                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                  handle ) );
-        PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                    PSA_KEY_DERIVATION_INPUT_INFO,
-                                                    input2->x, input2->len ) );
-    }
-    else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
-             PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
-    {
-        PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                    PSA_KEY_DERIVATION_INPUT_SEED,
-                                                    input1->x, input1->len ) );
-        PSA_ASSERT( psa_key_derivation_input_key( &operation,
-                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                  handle ) );
-        PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                    PSA_KEY_DERIVATION_INPUT_LABEL,
-                                                    input2->x, input2->len ) );
-    }
-    else
-    {
-        TEST_ASSERT( ! "Key derivation algorithm not supported" );
-    }
+    if( !setup_key_derivation_wrap( &operation, handle, alg,
+                                    input1->x, input1->len,
+                                    input2->x, input2->len,
+                                    requested_capacity ) )
+        goto exit;
 
     PSA_ASSERT( psa_key_derivation_get_capacity( &operation,
                                                  &current_capacity ) );