diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 9c303cb..1e7aaa8 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -3229,9 +3229,12 @@
  * Refer to the documentation of each key derivation or key agreement
  * algorithm for information.
  *
- * This function passes direct inputs. Some inputs must be passed as keys
- * using psa_key_derivation_input_key() instead of this function. Refer to
- * the documentation of individual step types for information.
+ * This function passes direct inputs, which is usually correct for
+ * non-secret inputs. To pass a secret input, which should be in a key
+ * object, call psa_key_derivation_input_key() instead of this function.
+ * Refer to the documentation of individual step types
+ * (`PSA_KEY_DERIVATION_INPUT_xxx` values of type ::psa_key_derivation_step_t)
+ * for more information.
  *
  * If this function returns an error status, the operation enters an error
  * state and must be aborted by calling psa_key_derivation_abort().
@@ -3274,10 +3277,13 @@
  * Refer to the documentation of each key derivation or key agreement
  * algorithm for information.
  *
- * This function passes key inputs. Some inputs must be passed as keys
- * of the appropriate type using this function, while others must be
- * passed as direct inputs using psa_key_derivation_input_bytes(). Refer to
- * the documentation of individual step types for information.
+ * This function obtains input from a key object, which is usually correct for
+ * secret inputs or for non-secret personalization strings kept in the key
+ * store. To pass a non-secret parameter which is not in the key store,
+ * call psa_key_derivation_input_bytes() instead of this function.
+ * Refer to the documentation of individual step types
+ * (`PSA_KEY_DERIVATION_INPUT_xxx` values of type ::psa_key_derivation_step_t)
+ * for more information.
  *
  * If this function returns an error status, the operation enters an error
  * state and must be aborted by calling psa_key_derivation_abort().
@@ -3298,7 +3304,8 @@
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         \c step is not compatible with the operation's algorithm.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \c step does not allow key inputs.
+ *         \c step does not allow key inputs of the given type
+ *         or does not allow key inputs at all.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
@@ -3368,6 +3375,8 @@
  *         \c private_key.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \c alg is not supported or is not a key derivation algorithm.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step does not allow an input resulting from a key agreement.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
@@ -3518,6 +3527,11 @@
  * In all cases, the data that is read is discarded from the operation.
  * The operation's capacity is decreased by the number of bytes read.
  *
+ * For algorithms that take an input step #PSA_KEY_DERIVATION_INPUT_SECRET,
+ * the input to that step must be provided with psa_key_derivation_input_key().
+ * Future versions of this specification may include additional restrictions
+ * on the derived key based on the attributes and strength of the secret key.
+ *
  * \param[in] attributes    The attributes for the new key.
  * \param[in,out] operation The key derivation operation object to read from.
  * \param[out] handle       On success, a handle to the newly created key.
@@ -3540,6 +3554,9 @@
  *         implementation in general or in this particular location.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         The provided key attributes are not valid for the operation.
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ *         The #PSA_KEY_DERIVATION_INPUT_SECRET input was not provided through
+ *         a key.
  * \retval #PSA_ERROR_BAD_STATE
  *         The operation state is not valid (it must be active and completed
  *         all required input steps).
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index f177d5d..9f55484 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -255,6 +255,7 @@
 struct psa_key_derivation_s
 {
     psa_algorithm_t alg;
+    unsigned int can_output_key : 1;
     size_t capacity;
     union
     {
@@ -268,7 +269,7 @@
 };
 
 /* This only zeroes out the first byte in the union, the rest is unspecified. */
-#define PSA_KEY_DERIVATION_OPERATION_INIT {0, 0, {0}}
+#define PSA_KEY_DERIVATION_OPERATION_INIT {0, 0, 0, {0}}
 static inline struct psa_key_derivation_s psa_key_derivation_operation_init( void )
 {
     const struct psa_key_derivation_s v = PSA_KEY_DERIVATION_OPERATION_INIT;
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index b3e0940..6b6a9f8 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -1618,31 +1618,43 @@
 
 /** A secret input for key derivation.
  *
- * This must be a key of type #PSA_KEY_TYPE_DERIVE.
+ * This should be a key of type #PSA_KEY_TYPE_DERIVE
+ * (passed to psa_key_derivation_input_key())
+ * or the shared secret resulting from a key agreement
+ * (obtained via psa_key_derivation_key_agreement()).
+ *
+ * The secret can also be a direct input (passed to
+ * key_derivation_input_bytes()). In this case, the derivation operation
+ * may not be used to derive keys: the operation will only allow
+ * psa_key_derivation_output_bytes(), not psa_key_derivation_output_key().
  */
 #define PSA_KEY_DERIVATION_INPUT_SECRET     ((psa_key_derivation_step_t)0x0101)
 
 /** A label for key derivation.
  *
- * This must be a direct input.
+ * This should be a direct input.
+ * It can also be a key of type #PSA_KEY_TYPE_RAW_DATA.
  */
 #define PSA_KEY_DERIVATION_INPUT_LABEL      ((psa_key_derivation_step_t)0x0201)
 
 /** A salt for key derivation.
  *
- * This must be a direct input.
+ * This should be a direct input.
+ * It can also be a key of type #PSA_KEY_TYPE_RAW_DATA.
  */
 #define PSA_KEY_DERIVATION_INPUT_SALT       ((psa_key_derivation_step_t)0x0202)
 
 /** An information string for key derivation.
  *
- * This must be a direct input.
+ * This should be a direct input.
+ * It can also be a key of type #PSA_KEY_TYPE_RAW_DATA.
  */
 #define PSA_KEY_DERIVATION_INPUT_INFO       ((psa_key_derivation_step_t)0x0203)
 
 /** A seed for key derivation.
  *
- * This must be a direct input.
+ * This should be a direct input.
+ * It can also be a key of type #PSA_KEY_TYPE_RAW_DATA.
  */
 #define PSA_KEY_DERIVATION_INPUT_SEED       ((psa_key_derivation_step_t)0x0204)
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index fe737d2..b9ea00f 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4787,6 +4787,9 @@
     if( psa_get_key_bits( attributes ) == 0 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    if( ! operation->can_output_key )
+        return( PSA_ERROR_NOT_PERMITTED );
+
     status = psa_start_key_creation( PSA_KEY_CREATION_DERIVE,
                                      attributes, handle, &slot, &driver );
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
@@ -5076,15 +5079,54 @@
 }
 #endif /* MBEDTLS_MD_C */
 
+/** Check whether the given key type is acceptable for the given
+ * input step of a key derivation.
+ *
+ * Secret inputs must have the type #PSA_KEY_TYPE_DERIVE.
+ * Non-secret inputs must have the type #PSA_KEY_TYPE_RAW_DATA.
+ * Both secret and non-secret inputs can alternatively have the type
+ * #PSA_KEY_TYPE_NONE, which is never the type of a key object, meaning
+ * that the input was passed as a buffer rather than via a key object.
+ */
+static int psa_key_derivation_check_input_type(
+    psa_key_derivation_step_t step,
+    psa_key_type_t key_type )
+{
+    switch( step )
+    {
+        case PSA_KEY_DERIVATION_INPUT_SECRET:
+            if( key_type == PSA_KEY_TYPE_DERIVE )
+                return( PSA_SUCCESS );
+            if( key_type == PSA_KEY_TYPE_NONE )
+                return( PSA_SUCCESS );
+            break;
+        case PSA_KEY_DERIVATION_INPUT_LABEL:
+        case PSA_KEY_DERIVATION_INPUT_SALT:
+        case PSA_KEY_DERIVATION_INPUT_INFO:
+        case PSA_KEY_DERIVATION_INPUT_SEED:
+            if( key_type == PSA_KEY_TYPE_RAW_DATA )
+                return( PSA_SUCCESS );
+            if( key_type == PSA_KEY_TYPE_NONE )
+                return( PSA_SUCCESS );
+            break;
+    }
+    return( PSA_ERROR_INVALID_ARGUMENT );
+}
+
 static psa_status_t psa_key_derivation_input_internal(
     psa_key_derivation_operation_t *operation,
     psa_key_derivation_step_t step,
+    psa_key_type_t key_type,
     const uint8_t *data,
     size_t data_length )
 {
     psa_status_t status;
     psa_algorithm_t kdf_alg = psa_key_derivation_get_kdf_alg( operation );
 
+    status = psa_key_derivation_check_input_type( step, key_type );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
 #if defined(MBEDTLS_MD_C)
     if( PSA_ALG_IS_HKDF( kdf_alg ) )
     {
@@ -5111,6 +5153,7 @@
         return( PSA_ERROR_BAD_STATE );
     }
 
+exit:
     if( status != PSA_SUCCESS )
         psa_key_derivation_abort( operation );
     return( status );
@@ -5122,10 +5165,8 @@
     const uint8_t *data,
     size_t data_length )
 {
-    if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
     return( psa_key_derivation_input_internal( operation, step,
+                                               PSA_KEY_TYPE_NONE,
                                                data, data_length ) );
 }
 
@@ -5136,23 +5177,23 @@
 {
     psa_key_slot_t *slot;
     psa_status_t status;
+
     status = psa_get_transparent_key( handle, &slot,
                                       PSA_KEY_USAGE_DERIVE,
                                       operation->alg );
     if( status != PSA_SUCCESS )
+    {
+        psa_key_derivation_abort( operation );
         return( status );
-    if( slot->attr.type != PSA_KEY_TYPE_DERIVE )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    /* Don't allow a key to be used as an input that is usually public.
-     * This is debatable. It's ok from a cryptographic perspective to
-     * use secret material as an input that is usually public. However
-     * the material should be dedicated to a particular input step,
-     * otherwise this may allow the key to be used in an unintended way
-     * and leak values derived from the key. So be conservative. */
-    if( step != PSA_KEY_DERIVATION_INPUT_SECRET )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
+    /* Passing a key object as a SECRET input unlocks the permission
+     * to output to a key object. */
+    if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
+        operation->can_output_key = 1;
+
     return( psa_key_derivation_input_internal( operation,
-                                               step,
+                                               step, slot->attr.type,
                                                slot->data.raw.data,
                                                slot->data.raw.bytes ) );
 }
@@ -5265,8 +5306,10 @@
         goto exit;
 
     /* Step 2: set up the key derivation to generate key material from
-     * the shared secret. */
+     * the shared secret. A shared secret is permitted wherever a key
+     * of type DERIVE is permitted. */
     status = psa_key_derivation_input_internal( operation, step,
+                                                PSA_KEY_TYPE_DERIVE,
                                                 shared_secret,
                                                 shared_secret_length );
 
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 8fe4cd2..6efdc01 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1864,69 +1864,145 @@
 Crypto derivation operation object initializers zero properly
 key_derivation_init:
 
-PSA key derivation: HKDF-SHA-256, good case [#1]
+PSA key derivation setup: HKDF-SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_SUCCESS
 
-PSA key derivation: HKDF-SHA-512, good case [#1]
+PSA key derivation setup: HKDF-SHA-512, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA512_C
 derive_setup:PSA_ALG_HKDF(PSA_ALG_SHA_512):PSA_SUCCESS
 
-PSA key derivation: TLS 1.2 PRF SHA-256, good case [#1]
+PSA key derivation setup: TLS 1.2 PRF SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_SUCCESS
 
-PSA key derivation: not a key derivation algorithm (HMAC)
+PSA key derivation setup: not a key derivation algorithm (HMAC)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_ALG_HMAC(PSA_ALG_SHA_256):PSA_ERROR_INVALID_ARGUMENT
 
-PSA key derivation: unsupported key derivation algorithm [#1]
+PSA key derivation setup: algorithm from bad hash
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_ALG_HKDF(PSA_ALG_CATEGORY_HASH):PSA_ERROR_NOT_SUPPORTED
 
-PSA key derivation: unsupported key derivation algorithm [#2]
+PSA key derivation setup: bad algorithm
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_ALG_CATEGORY_KEY_DERIVATION:PSA_ERROR_NOT_SUPPORTED
 
-PSA key derivation: HKDF-SHA-256, good case [#2]
+PSA key derivation: HKDF-SHA-256, good case, direct output
 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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
 
-PSA key derivation: HKDF-SHA-512, good case [#2]
+PSA key derivation: HKDF-SHA-256, good case, key output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
-PSA key derivation: TLS 1.2 PRF SHA-256, good case [#2]
+PSA key derivation: HKDF-SHA-256, bad key type, key output
 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
+# Whether we get NOT_PERMITTED or BAD_STATE for the output is an implementation
+# detail.
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_RAW_DATA:PSA_ERROR_NOT_PERMITTED
+
+PSA key derivation: HKDF-SHA-256, direct secret, direct output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
+
+PSA key derivation: HKDF-SHA-256, direct empty secret, direct output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
+
+PSA key derivation: HKDF-SHA-256, direct secret, key output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_RAW_DATA:PSA_ERROR_NOT_PERMITTED
+
+PSA key derivation: HKDF-SHA-256, direct empty secret, key output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_RAW_DATA:PSA_ERROR_NOT_PERMITTED
+
+PSA key derivation: HKDF-SHA-256, RAW_DATA key as salt
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_RAW_DATA:"412073616c74":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE:PSA_SUCCESS
+
+PSA key derivation: HKDF-SHA-256, RAW_DATA key as info
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_RAW_DATA:"4120696e666f":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE:PSA_SUCCESS
+
+PSA key derivation: HKDF-SHA-256, DERIVE key as salt, direct output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_DERIVE:"412073616c74":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: HKDF-SHA-256, DERIVE key as salt, key output
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+# Whether we get NOT_PERMITTED or BAD_STATE for the output is an implementation
+# detail.
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_DERIVE:"412073616c74":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_RAW_DATA:PSA_ERROR_BAD_STATE
+
+PSA key derivation: HKDF-SHA-256, DERIVE key as info
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_DERIVE:"4120696e666f":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+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_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE: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_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: TLS 1.2 PRF SHA-256, direct secret
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
+
+PSA key derivation: TLS 1.2 PRF SHA-256, direct empty secret
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
+
+PSA key derivation: TLS 1.2 PRF SHA-256, RAW_DATA key as seed
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_RAW_DATA:"612073656564":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE:PSA_SUCCESS
+
+PSA key derivation: TLS 1.2 PRF SHA-256, RAW_DATA key as label
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_RAW_DATA:"61206c6162656c":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE:PSA_SUCCESS
+
+PSA key derivation: TLS 1.2 PRF SHA-256, DERIVE key as seed
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_DERIVE:"612073656564":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: TLS 1.2 PRF SHA-256, DERIVE key as label
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_DERIVE:"61206c6162656c":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_TYPE_NONE: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:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE: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 +2148,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..87529ac 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4273,17 +4273,17 @@
 
 /* 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 expected_status_arg1,
+                   int step_arg2, int key_type_arg2, data_t *input2,
                    int expected_status_arg2,
-                   int expected_status_arg3 )
+                   int step_arg3, int key_type_arg3, data_t *input3,
+                   int expected_status_arg3,
+                   int output_key_type_arg, int expected_output_status_arg )
 {
     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};
@@ -4292,40 +4292,62 @@
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     size_t i;
+    psa_key_type_t output_key_type = output_key_type_arg;
+    psa_key_handle_t output_handle = 0;
+    psa_status_t expected_output_status = expected_output_status_arg;
+    psa_status_t actual_output_status;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
     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] != PSA_KEY_TYPE_NONE )
         {
-            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] );
         }
     }
 
+    if( output_key_type != PSA_KEY_TYPE_NONE )
+    {
+        psa_reset_key_attributes( &attributes );
+        psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
+        psa_set_key_bits( &attributes, 8 );
+        actual_output_status =
+            psa_key_derivation_output_key( &attributes, &operation,
+                                           &output_handle );
+    }
+    else
+    {
+        uint8_t buffer[1];
+        actual_output_status =
+            psa_key_derivation_output_bytes( &operation,
+                                             buffer, sizeof( buffer ) );
+    }
+    TEST_EQUAL( actual_output_status, expected_output_status );
+
 exit:
     psa_key_derivation_abort( &operation );
     for( i = 0; i < ARRAY_LENGTH( handles ); i++ )
         psa_destroy_key( handles[i] );
+    psa_destroy_key( output_handle );
     PSA_DONE( );
 }
 /* END_CASE */