Convert derive_output to the new KDF API
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 858356d..8e638b6 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4213,15 +4213,17 @@
 
 /* BEGIN_CASE */
 void derive_output( int alg_arg,
-                    data_t *key_data,
-                    data_t *salt,
-                    data_t *label,
+                    int step1_arg, data_t *input1,
+                    int step2_arg, data_t *input2,
+                    int step3_arg, data_t *input3,
                     int requested_capacity_arg,
                     data_t *expected_output1,
                     data_t *expected_output2 )
 {
-    psa_key_handle_t handle = 0;
     psa_algorithm_t alg = alg_arg;
+    psa_key_derivation_step_t steps[] = {step1_arg, step2_arg, step3_arg};
+    data_t *inputs[] = {input1, input2, input3};
+    psa_key_handle_t handles[] = {0, 0, 0};
     size_t requested_capacity = requested_capacity_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     uint8_t *expected_outputs[2] =
@@ -4234,7 +4236,7 @@
     size_t current_capacity;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status;
-    unsigned i;
+    size_t i;
 
     for( i = 0; i < ARRAY_LENGTH( expected_outputs ); i++ )
     {
@@ -4250,35 +4252,32 @@
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
 
-    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
-                                &handle ) );
-
     /* Extraction phase. */
-    if( PSA_ALG_IS_HKDF( alg ) )
+    PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
+    PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
+                                                 requested_capacity ) );
+    for( i = 0; i < ARRAY_LENGTH( steps ); i++ )
     {
-        PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-        PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
-                                                     requested_capacity ) );
-        PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                    PSA_KEY_DERIVATION_INPUT_SALT,
-                                                    salt->x, salt->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,
-                                                    label->x, label->len ) );
+        switch( steps[i] )
+        {
+            case 0:
+                break;
+            case PSA_KEY_DERIVATION_INPUT_SECRET:
+                PSA_ASSERT( psa_import_key( &attributes,
+                                            inputs[i]->x, inputs[i]->len,
+                                            &handles[i] ) );
+                PSA_ASSERT( psa_key_derivation_input_key(
+                                &operation, steps[i],
+                                handles[i] ) );
+                break;
+            default:
+                PSA_ASSERT( psa_key_derivation_input_bytes(
+                                &operation, steps[i],
+                                inputs[i]->x, inputs[i]->len ) );
+                break;
+        }
     }
-#if defined(PSA_PRE_1_0_KEY_DERIVATION)
-    else
-    {
-        // legacy
-        PSA_ASSERT( psa_key_derivation( &operation, handle, alg,
-                                        salt->x, salt->len,
-                                        label->x, label->len,
-                                        requested_capacity ) );
-    }
-#endif
+
     PSA_ASSERT( psa_key_derivation_get_capacity( &operation,
                                                  &current_capacity ) );
     TEST_EQUAL( current_capacity, requested_capacity );
@@ -4321,7 +4320,8 @@
 exit:
     mbedtls_free( output_buffer );
     psa_key_derivation_abort( &operation );
-    psa_destroy_key( handle );
+    for( i = 0; i < ARRAY_LENGTH( handles ); i++ )
+        psa_destroy_key( handles[i] );
     PSA_DONE( );
 }
 /* END_CASE */