diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index d713ac823..7591b6bb1 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4763,6 +4763,9 @@ psa_status_t psa_key_derivation_output_key( const psa_key_attributes_t *attribut
     if( psa_get_key_bits( attributes ) == 0 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    if( operation->alg == PSA_ALG_NONE )
+        return( PSA_ERROR_BAD_STATE );
+
     if( ! operation->can_output_key )
         return( PSA_ERROR_NOT_PERMITTED );
 
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 6dbf18c8e..9060159f5 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -4262,7 +4262,7 @@ PSA key derivation: HKDF-SHA-256, bad key type, key output
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 # 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
+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_BAD_STATE
 
 PSA key derivation: HKDF-SHA-256, direct secret, direct output
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 1aa95275a..d28de0c20 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -5934,7 +5934,7 @@ void derive_input( int alg_arg,
     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_type( &attributes, output_key_type );
         psa_set_key_bits( &attributes, 8 );
         actual_output_status =
             psa_key_derivation_output_key( &attributes, &operation,