diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index c9e681746..e58abf9ca 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1757,39 +1757,31 @@ key_derivation_init:
 
 PSA key derivation: HKDF-SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HKDF(PSA_ALG_SHA_256):"":"":42:PSA_SUCCESS
+derive_setup:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_SUCCESS
 
 PSA key derivation: HKDF-SHA-512, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA512_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HKDF(PSA_ALG_SHA_512):"":"":42:PSA_SUCCESS
-
-PSA key derivation: HKDF-SHA-256, bad key type
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HKDF(PSA_ALG_SHA_256):"":"":42:PSA_ERROR_INVALID_ARGUMENT
+derive_setup:PSA_ALG_HKDF(PSA_ALG_SHA_512):PSA_SUCCESS
 
 PSA key derivation: TLS 1.2 PRF SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"":"":42:PSA_SUCCESS
-
-PSA key derivation: TLS 1.2 PRF SHA-256, bad key type
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_RAW_DATA:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"":"":42:PSA_ERROR_INVALID_ARGUMENT
+derive_setup:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_SUCCESS
 
 PSA key derivation: not a key derivation algorithm (selection)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_SELECT_RAW:"":"":42:PSA_ERROR_INVALID_ARGUMENT
+derive_setup:PSA_ALG_SELECT_RAW:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: not a key derivation algorithm (HMAC)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_256):"":"":42:PSA_ERROR_INVALID_ARGUMENT
+derive_setup:PSA_ALG_HMAC(PSA_ALG_SHA_256):PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: unsupported key derivation algorithm
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HKDF(PSA_ALG_CATEGORY_HASH):"":"":42:PSA_ERROR_NOT_SUPPORTED
+derive_setup::PSA_ALG_HKDF(PSA_ALG_CATEGORY_HASH):PSA_ERROR_NOT_SUPPORTED
 
 PSA key derivation: unsupported key derivation algorithm
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_CATEGORY_KEY_DERIVATION:"":"":42:PSA_ERROR_NOT_SUPPORTED
+derive_setup:PSA_ALG_CATEGORY_KEY_DERIVATION:PSA_ERROR_NOT_SUPPORTED
 
 PSA key derivation: HKDF-SHA-256, good case
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
@@ -1963,17 +1955,9 @@ PSA key derivation: HKDF SHA-1, request maximum capacity
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA1_C
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_1):"0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c":"":"":255 * 20:"2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48":""
 
-PSA key derivation: HKDF SHA-256, request too much capacity
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HKDF(PSA_ALG_SHA_256):"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32 + 1:PSA_ERROR_INVALID_ARGUMENT
-
-PSA key derivation: HKDF SHA-1, request too much capacity
-depends_on:MBEDTLS_MD_C:MBEDTLS_SHA1_C
-derive_setup:PSA_KEY_TYPE_DERIVE:"0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c":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_setup:PSA_KEY_TYPE_DERIVE:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"":"":100:PSA_ERROR_INVALID_ARGUMENT
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+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
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index a049ee8cb..316f2edcd 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4034,41 +4034,20 @@ void key_derivation_init( )
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:PSA_PRE_1_0_KEY_DERIVATION */
-void derive_setup( int key_type_arg,
-                   data_t *key_data,
-                   int alg_arg,
-                   data_t *salt,
-                   data_t *label,
-                   int requested_capacity_arg,
-                   int expected_status_arg )
+/* BEGIN_CASE */
+void derive_setup( int alg_arg, int expected_status_arg )
 {
-    psa_key_handle_t handle = 0;
-    size_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
-    size_t requested_capacity = requested_capacity_arg;
     psa_status_t expected_status = expected_status_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     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_import_key( &attributes, key_data->x, key_data->len,
-                                &handle ) );
-
-    TEST_EQUAL( psa_key_derivation( &operation, handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    requested_capacity ),
+    TEST_EQUAL( psa_key_derivation_setup( &operation, alg ),
                 expected_status );
 
 exit:
     psa_key_derivation_abort( &operation );
-    psa_destroy_key( handle );
     PSA_DONE( );
 }
 /* END_CASE */