diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index a0e7f7a..53aa41a 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -52,6 +52,12 @@
 PSA invalid handle (largest plausible handle)
 invalid_handle:-1
 
+PSA import: invalid type (0)
+import:"0123":PSA_KEY_TYPE_NONE:0:PSA_ERROR_NOT_SUPPORTED
+
+PSA import: invalid type (PSA_KEY_TYPE_CATEGORY_MASK)
+import:"0123":PSA_KEY_TYPE_CATEGORY_MASK:0:PSA_ERROR_NOT_SUPPORTED
+
 PSA import AES: bad key size
 depends_on:MBEDTLS_AES_C
 import:"0123456789abcdef":PSA_KEY_TYPE_AES:0:PSA_ERROR_INVALID_ARGUMENT
@@ -2083,16 +2089,24 @@
 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 derivation: invalid type (0)
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_NOT_SUPPORTED
+
+PSA key derivation: invalid type (PSA_KEY_TYPE_CATEGORY_MASK)
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_ERROR_NOT_SUPPORTED
+
 # This test assumes that PSA_MAX_KEY_BITS (currently 65536-8 bits = 8191 bytes
 # and not expected to be raised any time soon) is less than the maximum
 # output from HKDF-SHA512 (255*64 = 16320 bytes).
 PSA key derivation: largest possible key
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA512_C:!PSA_PRE_1_0_KEY_DERIVATION
-derive_large_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_MAX_KEY_BITS:PSA_SUCCESS
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS:PSA_SUCCESS
 
 PSA key derivation: key too large
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA512_C:!PSA_PRE_1_0_KEY_DERIVATION
-derive_large_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED
 
 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
@@ -2192,6 +2206,9 @@
 PSA generate random: 260 bytes
 generate_random:260
 
+PSA generate key: bad type (0)
+generate_key:PSA_KEY_TYPE_NONE:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+
 PSA generate key: bad type (PSA_KEY_TYPE_CATEGORY_MASK)
 generate_key:PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
 
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 8ed7a7d..60514f8 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4620,14 +4620,15 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void derive_large_key( int alg_arg,
-                       data_t *key_data, data_t *input1, data_t *input2,
-                       int bits_arg,
-                       int expected_status_arg )
+void derive_key( int alg_arg,
+                 data_t *key_data, data_t *input1, data_t *input2,
+                 int type_arg, int bits_arg,
+                 int expected_status_arg )
 {
     psa_key_handle_t base_handle = 0;
     psa_key_handle_t derived_handle = 0;
     psa_algorithm_t alg = alg_arg;
+    psa_key_type_t type = type_arg;
     size_t bits = bits_arg;
     psa_status_t expected_status = expected_status_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
@@ -4649,7 +4650,7 @@
 
     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 );
+    psa_set_key_type( &derived_attributes, type );
     psa_set_key_bits( &derived_attributes, bits );
     TEST_EQUAL( psa_key_derivation_output_key( &derived_attributes, &operation,
                                                &derived_handle ),