PSA crypto KDF: test bytes/key input independently of the step type

This commit only makes derive_input more flexible so that the key
derivation API can be tested with different key types and raw data for
each input step. The behavior of the test cases remains the same.
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 3225bef..79ef9a8 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4273,17 +4273,16 @@
 
 /* BEGIN_CASE */
 void derive_input( int alg_arg,
-                   int key_type_arg,
-                   int step1_arg, data_t *input1,
-                   int step2_arg, data_t *input2,
-                   int step3_arg, data_t *input3,
+                   int step_arg1, int key_type_arg1, data_t *input1,
+                   int step_arg2, int key_type_arg2, data_t *input2,
+                   int step_arg3, int key_type_arg3, data_t *input3,
                    int expected_status_arg1,
                    int expected_status_arg2,
                    int expected_status_arg3 )
 {
     psa_algorithm_t alg = alg_arg;
-    size_t key_type = key_type_arg;
-    psa_key_derivation_step_t steps[] = {step1_arg, step2_arg, step3_arg};
+    psa_key_derivation_step_t steps[] = {step_arg1, step_arg2, step_arg3};
+    psa_key_type_t key_types[] = {key_type_arg1, key_type_arg2, key_type_arg3};
     psa_status_t expected_statuses[] = {expected_status_arg1,
                                         expected_status_arg2,
                                         expected_status_arg3};
@@ -4297,28 +4296,27 @@
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
     psa_set_key_algorithm( &attributes, alg );
-    psa_set_key_type( &attributes, key_type );
 
     PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
 
     for( i = 0; i < ARRAY_LENGTH( steps ); i++ )
     {
-        switch( steps[i] )
+        if( key_types[i] != 0 )
         {
-            case PSA_KEY_DERIVATION_INPUT_SECRET:
-                PSA_ASSERT( psa_import_key( &attributes,
-                                            inputs[i]->x, inputs[i]->len,
-                                            &handles[i] ) );
-                TEST_EQUAL( psa_key_derivation_input_key( &operation, steps[i],
-                                                          handles[i] ),
-                            expected_statuses[i] );
-                break;
-            default:
-                TEST_EQUAL( psa_key_derivation_input_bytes(
-                                    &operation, steps[i],
-                                    inputs[i]->x, inputs[i]->len ),
-                            expected_statuses[i] );
-                break;
+            psa_set_key_type( &attributes, key_types[i] );
+            PSA_ASSERT( psa_import_key( &attributes,
+                                        inputs[i]->x, inputs[i]->len,
+                                        &handles[i] ) );
+            TEST_EQUAL( psa_key_derivation_input_key( &operation, steps[i],
+                                                      handles[i] ),
+                        expected_statuses[i] );
+        }
+        else
+        {
+            TEST_EQUAL( psa_key_derivation_input_bytes(
+                            &operation, steps[i],
+                            inputs[i]->x, inputs[i]->len ),
+                        expected_statuses[i] );
         }
     }