diff --git a/programs/psa/key_ladder_demo.c b/programs/psa/key_ladder_demo.c
index 426e41f..91e5178 100644
--- a/programs/psa/key_ladder_demo.c
+++ b/programs/psa/key_ladder_demo.c
@@ -63,27 +63,25 @@
 
 #include "mbedtls/platform_util.h" // for mbedtls_platform_zeroize
 
+#include <psa/crypto.h>
+
 /* If the build options we need are not enabled, compile a placeholder. */
 #if !defined(MBEDTLS_SHA256_C) || !defined(MBEDTLS_MD_C) ||     \
     !defined(MBEDTLS_AES_C) || !defined(MBEDTLS_CCM_C) ||       \
     !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(MBEDTLS_FS_IO) ||\
-    !defined(PSA_PRE_1_0_KEY_DERIVATION)
+    defined(PSA_PRE_1_0_KEY_DERIVATION)
 int main( void )
 {
     printf("MBEDTLS_SHA256_C and/or MBEDTLS_MD_C and/or "
            "MBEDTLS_AES_C and/or MBEDTLS_CCM_C and/or "
            "MBEDTLS_PSA_CRYPTO_C and/or MBEDTLS_FS_IO and/or "
-           "PSA_PRE_1_0_KEY_DERIVATION not defined.\n");
+           "not defined and/or PSA_PRE_1_0_KEY_DERIVATION defined.\n");
     return( 0 );
 }
 #else
 
 /* The real program starts here. */
 
-
-
-#include <psa/crypto.h>
-
 /* Run a system function and bail out if it fails. */
 #define SYS_CHECK( expr )                                       \
     do                                                          \
@@ -281,7 +279,7 @@
 {
     psa_status_t status = PSA_SUCCESS;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_derivation_operation_t generator = PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     size_t i;
 
     psa_set_key_usage_flags( &attributes,
@@ -295,26 +293,28 @@
     {
         /* Start deriving material from the master key (if i=0) or from
          * the current intermediate key (if i>0). */
-        PSA_CHECK( psa_key_derivation(
-                       &generator,
-                       *key_handle,
-                       KDF_ALG,
-                       DERIVE_KEY_SALT, DERIVE_KEY_SALT_LENGTH,
-                       (uint8_t*) ladder[i], strlen( ladder[i] ),
-                       KEY_SIZE_BYTES ) );
+        PSA_CHECK( psa_key_derivation_setup( &operation, KDF_ALG ) );
+        PSA_CHECK( psa_key_derivation_input_bytes(
+                       &operation, PSA_KEY_DERIVATION_INPUT_SALT,
+                       DERIVE_KEY_SALT, DERIVE_KEY_SALT_LENGTH ) );
+        PSA_CHECK( psa_key_derivation_input_key(
+                       &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
+                       *key_handle ) );
+        PSA_CHECK( psa_key_derivation_input_bytes(
+                       &operation, PSA_KEY_DERIVATION_INPUT_INFO,
+                       (uint8_t*) ladder[i], strlen( ladder[i] ) ) );
         /* When the parent key is not the master key, destroy it,
          * since it is no longer needed. */
         PSA_CHECK( psa_close_key( *key_handle ) );
         *key_handle = 0;
-        /* Use the generator obtained from the parent key to create
-         * the next intermediate key. */
-        PSA_CHECK( psa_key_derivation_output_key( &attributes, &generator,
-                                             key_handle ) );
-        PSA_CHECK( psa_key_derivation_abort( &generator ) );
+        /* Derive the next intermediate key from the parent key. */
+        PSA_CHECK( psa_key_derivation_output_key( &attributes, &operation,
+                                                  key_handle ) );
+        PSA_CHECK( psa_key_derivation_abort( &operation ) );
     }
 
 exit:
-    psa_key_derivation_abort( &generator );
+    psa_key_derivation_abort( &operation );
     if( status != PSA_SUCCESS )
     {
         psa_close_key( *key_handle );
@@ -330,31 +330,33 @@
 {
     psa_status_t status = PSA_SUCCESS;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_derivation_operation_t generator = PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
 
     *wrapping_key_handle = 0;
+
+    /* Set up a key derivation operation from the key derived from
+     * the master key. */
+    PSA_CHECK( psa_key_derivation_setup( &operation, KDF_ALG ) );
+    PSA_CHECK( psa_key_derivation_input_bytes(
+                   &operation, PSA_KEY_DERIVATION_INPUT_SALT,
+                   WRAPPING_KEY_SALT, WRAPPING_KEY_SALT_LENGTH ) );
+    PSA_CHECK( psa_key_derivation_input_key(
+                   &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
+                   derived_key_handle ) );
+    PSA_CHECK( psa_key_derivation_input_bytes(
+                   &operation, PSA_KEY_DERIVATION_INPUT_INFO,
+                   NULL, 0 ) );
+
+    /* Create the wrapping key. */
     psa_set_key_usage_flags( &attributes, usage );
     psa_set_key_algorithm( &attributes, WRAPPING_ALG );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
     psa_set_key_bits( &attributes, WRAPPING_KEY_BITS );
-
-    PSA_CHECK( psa_key_derivation(
-                   &generator,
-                   derived_key_handle,
-                   KDF_ALG,
-                   WRAPPING_KEY_SALT, WRAPPING_KEY_SALT_LENGTH,
-                   NULL, 0,
-                   PSA_BITS_TO_BYTES( WRAPPING_KEY_BITS ) ) );
-    PSA_CHECK( psa_key_derivation_output_key( &attributes, &generator,
-                                         wrapping_key_handle ) );
+    PSA_CHECK( psa_key_derivation_output_key( &attributes, &operation,
+                                              wrapping_key_handle ) );
 
 exit:
-    psa_key_derivation_abort( &generator );
-    if( status != PSA_SUCCESS )
-    {
-        psa_close_key( *wrapping_key_handle );
-        *wrapping_key_handle = 0;
-    }
+    psa_key_derivation_abort( &operation );
     return( status );
 }
 
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index f618e13..53f8422 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1466,6 +1466,14 @@
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
 import_and_exercise_key:"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):256:PSA_ALG_ECDH
 
+PSA import/exercise: HKDF SHA-256
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+import_and_exercise_key:"c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0":PSA_KEY_TYPE_DERIVE:192:PSA_ALG_HKDF(PSA_ALG_SHA_256)
+
+PSA import/exercise: TLS 1.2 PRF SHA-256
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+import_and_exercise_key:"c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0":PSA_KEY_TYPE_DERIVE:192:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256)
+
 PSA sign: RSA PKCS#1 v1.5, raw
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15
 sign_deterministic:PSA_KEY_TYPE_RSA_KEY_PAIR:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_ALG_RSA_PKCS1V15_SIGN_RAW:"616263":"2c7744983f023ac7bb1c55529d83ed11a76a7898a1bb5ce191375a4aa7495a633d27879ff58eba5a57371c34feb1180e8b850d552476ebb5634df620261992f12ebee9097041dbbea85a42d45b344be5073ceb772ffc604954b9158ba81ec3dc4d9d65e3ab7aa318165f38c36f841f1c69cb1cfa494aa5cbb4d6c0efbafb043a"
@@ -1819,9 +1827,13 @@
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
 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
 
-PSA key derivation: invalid state (double generate + read past capacity)
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-test_derive_invalid_key_derivation_state:
+PSA key derivation: HKDF invalid state (double generate + read past capacity)
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+test_derive_invalid_key_derivation_state:PSA_ALG_HKDF(PSA_ALG_SHA_256)
+
+PSA key derivation: TLS 1.2 PRF invalid state (double generate + read past capacity)
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+test_derive_invalid_key_derivation_state:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256)
 
 PSA key derivation: invalid state (call read/get_capacity after init and abort)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
@@ -1987,42 +1999,82 @@
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32
 
+PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity minus 1
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32 - 1
+
+PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32
+
 PSA key derivation: HKDF SHA-256, exercise AES128-CTR
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CTR
 
 PSA key derivation: HKDF SHA-256, exercise AES256-CTR
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:256:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CTR
 
 PSA key derivation: HKDF SHA-256, exercise DES-CBC
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
 
 PSA key derivation: HKDF SHA-256, exercise 2-key 3DES-CBC
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
 
 PSA key derivation: HKDF SHA-256, exercise 3-key 3DES-CBC
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
 
 PSA key derivation: HKDF SHA-256, exercise HMAC-SHA-256
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_HMAC:256:PSA_KEY_USAGE_SIGN:PSA_ALG_HMAC(PSA_ALG_SHA_256)
 
-PSA key derivation: HKDF SHA-256, exercise HKDF-SHA-256
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DERIVE:400:PSA_KEY_USAGE_DERIVE:PSA_ALG_HKDF(PSA_ALG_SHA_256)
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise AES128-CTR
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CTR
 
-PSA key derivation: HKDF SHA-256, derive key, 16+32
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise AES256-CTR
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:256:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CTR
+
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise DES-CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
+
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise 2-key 3DES-CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
+
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise 3-key 3DES-CBC
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_PKCS7
+
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise HMAC-SHA-256
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_HMAC:256:PSA_KEY_USAGE_SIGN:PSA_ALG_HMAC(PSA_ALG_SHA_256)
+
+PSA key derivation: TLS 1.2 PRF SHA-256, exercise HKDF-SHA-256
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_exercise:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_DERIVE:400:PSA_KEY_USAGE_DERIVE:PSA_ALG_HKDF(PSA_ALG_SHA_256)
+
+PSA key derivation: HKDF SHA-256, derive key export, 16+32
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_export:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":16:32
 
-PSA key derivation: HKDF SHA-256, derive key, 1+41
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+PSA key derivation: HKDF SHA-256, derive key export, 1+41
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
 derive_key_export:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":1:41
 
+PSA key derivation: TLS 1.2 PRF SHA-256, derive key export, 16+32
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_export:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":16:32
+
+PSA key derivation: TLS 1.2 PRF SHA-256, derive key export, 1+41
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key_export:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":1:41
+
 PSA key agreement setup: ECDH + HKDF-SHA-256: good
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 key_agreement_setup:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 48f5337..1d06d62 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -517,48 +517,76 @@
     return( 0 );
 }
 
+static int setup_key_derivation_wrap( psa_key_derivation_operation_t* operation,
+                                      psa_key_handle_t handle,
+                                      psa_algorithm_t alg,
+                                      unsigned char* input1, size_t input1_length,
+                                      unsigned char* input2, size_t input2_length,
+                                      size_t capacity )
+{
+    PSA_ASSERT( psa_key_derivation_setup( operation, alg ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SALT,
+                                                    input1, input1_length ) );
+        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,
+                                                    input2,
+                                                    input2_length ) );
+    }
+    else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SEED,
+                                                    input1, input1_length ) );
+        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_LABEL,
+                                                    input2, input2_length ) );
+    }
+    else
+    {
+        TEST_ASSERT( ! "Key derivation algorithm not supported" );
+    }
+
+    PSA_ASSERT( psa_key_derivation_set_capacity( operation, capacity ) );
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+
 static int exercise_key_derivation_key( psa_key_handle_t handle,
                                         psa_key_usage_t usage,
                                         psa_algorithm_t alg )
 {
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    unsigned char label[16] = "This is a label.";
-    size_t label_length = sizeof( label );
-    unsigned char seed[16] = "abcdefghijklmnop";
-    size_t seed_length = sizeof( seed );
+    unsigned char input1[] = "Input 1";
+    size_t input1_length = sizeof( input1 );
+    unsigned char input2[] = "Input 2";
+    size_t input2_length = sizeof( input2 );
     unsigned char output[1];
+    size_t capacity = sizeof( output );
 
     if( usage & PSA_KEY_USAGE_DERIVE )
     {
-        if( PSA_ALG_IS_HKDF( alg ) )
-        {
-            PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-            PSA_ASSERT( psa_key_derivation_input_bytes( &operation,
-                                                        PSA_KEY_DERIVATION_INPUT_SALT,
-                                                        label,
-                                                        label_length ) );
-            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,
-                                                        seed,
-                                                        seed_length ) );
-        }
-#if defined(PSA_PRE_1_0_KEY_DERIVATION)
-        else
-        {
-            // legacy
-            PSA_ASSERT( psa_key_derivation( &operation,
-                                            handle, alg,
-                                            label, label_length,
-                                            seed, seed_length,
-                                            sizeof( output ) ) );
-        }
-#endif
+        if( !setup_key_derivation_wrap( &operation, handle, alg,
+                                        input1, input1_length,
+                                        input2, input2_length, capacity ) )
+            goto exit;
+
         PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
                                                      output,
-                                                     sizeof( output ) ) );
+                                                     capacity ) );
         PSA_ASSERT( psa_key_derivation_abort( &operation ) );
     }
 
@@ -4135,13 +4163,17 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:PSA_PRE_1_0_KEY_DERIVATION */
-void test_derive_invalid_key_derivation_state( )
+/* BEGIN_CASE */
+void test_derive_invalid_key_derivation_state( int alg_arg )
 {
+    psa_algorithm_t alg = alg_arg;
     psa_key_handle_t handle = 0;
     size_t key_type = PSA_KEY_TYPE_DERIVE;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    psa_algorithm_t alg = PSA_ALG_HKDF( PSA_ALG_SHA_256 );
+    unsigned char input1[] = "Input 1";
+    size_t input1_length = sizeof( input1 );
+    unsigned char input2[] = "Input 2";
+    size_t input2_length = sizeof( input2 );
     uint8_t buffer[42];
     size_t capacity = sizeof( buffer );
     const uint8_t key_data[22] = { 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b,
@@ -4160,16 +4192,14 @@
                                 &handle ) );
 
     /* valid key derivation */
-    PSA_ASSERT(  psa_key_derivation( &operation, handle, alg,
-                                     NULL, 0,
-                                     NULL, 0,
-                                     capacity ) );
+    if( !setup_key_derivation_wrap( &operation, handle, alg,
+                                    input1, input1_length,
+                                    input2, input2_length,
+                                    capacity ) )
+        goto exit;
 
     /* state of operation shouldn't allow additional generation */
-    TEST_EQUAL(  psa_key_derivation( &operation, handle, alg,
-                                     NULL, 0,
-                                     NULL, 0,
-                                     capacity ),
+    TEST_EQUAL(  psa_key_derivation_setup( &operation, alg ),
                  PSA_ERROR_BAD_STATE );
 
     PSA_ASSERT( psa_key_derivation_output_bytes( &operation, buffer, capacity ) );
@@ -4331,8 +4361,8 @@
 /* BEGIN_CASE */
 void derive_full( int alg_arg,
                   data_t *key_data,
-                  data_t *salt,
-                  data_t *label,
+                  data_t *input1,
+                  data_t *input2,
                   int requested_capacity_arg )
 {
     psa_key_handle_t handle = 0;
@@ -4353,33 +4383,12 @@
     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 ) );
-        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 ) );
-    }
+    if( !setup_key_derivation_wrap( &operation, handle, alg,
+                                    input1->x, input1->len,
+                                    input2->x, input2->len,
+                                    requested_capacity ) )
+        goto exit;
 
-#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, expected_capacity );
@@ -4412,11 +4421,11 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:PSA_PRE_1_0_KEY_DERIVATION */
+/* BEGIN_CASE */
 void derive_key_exercise( int alg_arg,
                           data_t *key_data,
-                          data_t *salt,
-                          data_t *label,
+                          data_t *input1,
+                          data_t *input2,
                           int derived_type_arg,
                           int derived_bits_arg,
                           int derived_usage_arg,
@@ -4443,10 +4452,11 @@
                                 &base_handle ) );
 
     /* Derive a key. */
-    PSA_ASSERT( psa_key_derivation( &operation, base_handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    capacity ) );
+    if ( setup_key_derivation_wrap( &operation, base_handle, alg,
+                                    input1->x, input1->len,
+                                    input2->x, input2->len, capacity ) )
+        goto exit;
+
     psa_set_key_usage_flags( &attributes, derived_usage );
     psa_set_key_algorithm( &attributes, derived_alg );
     psa_set_key_type( &attributes, derived_type );
@@ -4472,11 +4482,11 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:PSA_PRE_1_0_KEY_DERIVATION */
+/* BEGIN_CASE */
 void derive_key_export( int alg_arg,
                         data_t *key_data,
-                        data_t *salt,
-                        data_t *label,
+                        data_t *input1,
+                        data_t *input2,
                         int bytes1_arg,
                         int bytes2_arg )
 {
@@ -4504,20 +4514,22 @@
                                 &base_handle ) );
 
     /* Derive some material and output it. */
-    PSA_ASSERT( psa_key_derivation( &operation, base_handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    capacity ) );
+    if( !setup_key_derivation_wrap( &operation, base_handle, alg,
+                                    input1->x, input1->len,
+                                    input2->x, input2->len, capacity ) )
+        goto exit;
+
     PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
                                                  output_buffer,
                                                  capacity ) );
     PSA_ASSERT( psa_key_derivation_abort( &operation ) );
 
     /* Derive the same output again, but this time store it in key objects. */
-    PSA_ASSERT( psa_key_derivation( &operation, base_handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    capacity ) );
+    if( !setup_key_derivation_wrap( &operation, base_handle, alg,
+                                    input1->x, input1->len,
+                                    input2->x, input2->len, capacity ) )
+        goto exit;
+
     psa_set_key_usage_flags( &derived_attributes, PSA_KEY_USAGE_EXPORT );
     psa_set_key_algorithm( &derived_attributes, 0 );
     psa_set_key_type( &derived_attributes, PSA_KEY_TYPE_RAW_DATA );