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] );
         }
     }