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.data b/tests/suites/test_suite_psa_crypto.data
index 5866bae..cf95698 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1890,43 +1890,47 @@
 
 PSA key derivation: HKDF-SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SALT:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:0:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
 
 PSA key derivation: HKDF-SHA-512, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA512_C
-derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_512):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SALT:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_512):PSA_KEY_DERIVATION_INPUT_SALT:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:0:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
 
 PSA key derivation: HKDF-SHA-256, bad key type
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_TYPE_RAW_DATA:PSA_KEY_DERIVATION_INPUT_SALT:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_SUCCESS
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:0:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_SUCCESS
 
 PSA key derivation: TLS 1.2 PRF SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
 
 PSA key derivation: TLS 1.2 PRF SHA-256, key first
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS 1.2 PRF SHA-256, label first
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS 1.2 PRF SHA-256, early label
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS 1.2 PRF SHA-256, double seed
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS 1.2 PRF SHA-256, double key
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS 1.2 PRF SHA-256, bad key type
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_TYPE_RAW_DATA:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE
+
+PSA key derivation: TLS 1.2 PSK-to-MS, SHA-256, PSK too long (160 Bytes)
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:0:"":PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_KEY_DERIVATION_INPUT_LABEL:0:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE
 
 PSA key derivation: HKDF invalid state (double generate + read past capacity)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
@@ -2072,10 +2076,6 @@
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA1_C
 derive_set_capacity:PSA_ALG_HKDF(PSA_ALG_SHA_1):255 * 20 + 1:PSA_ERROR_INVALID_ARGUMENT
 
-PSA key derivation: TLS 1.2 PSK-to-MS, SHA-256, PSK too long (160 Bytes)
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_input:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE
-
 PSA key derivation: over capacity 42: output 42+1
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"ff"
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] );
         }
     }