diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3ecaebb7f..8c680778e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5284,6 +5284,11 @@ static psa_status_t psa_hkdf_input( psa_hkdf_key_derivation_t *hkdf,
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
             if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) )
                 return( PSA_ERROR_INVALID_ARGUMENT );
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXPAND)
+            if( PSA_ALG_IS_HKDF_EXPAND( kdf_alg ) &&
+                hkdf->state == HKDF_STATE_INIT )
+                return( PSA_ERROR_BAD_STATE );
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
             if( hkdf->state == HKDF_STATE_OUTPUT )
                 return( PSA_ERROR_BAD_STATE );
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index ed4fd2658..6f319938d 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5307,10 +5307,6 @@ PSA key derivation: HKDF-Expand SHA-1, RFC5869 #7, out 42+0 k
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_1
 derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_1):PSA_KEY_DERIVATION_INPUT_SECRET:"2adccada18779e7c2077ad2eb19d3f3e731385dd":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48":"":0:1:0
 
-PSA key derivation: HKDF-Expand SHA-256, RFC5869 #1 k info before secret
-depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
-derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"":0:1:0
-
 PSA key derivation: HKDF-Expand SHA-256, RFC5869 #1, out 42+0 b
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SECRET:"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"":0:0:0
@@ -5359,10 +5355,6 @@ PSA key derivation: HKDF-Expand SHA-1, RFC5869 #7, out 42+0 b
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_1
 derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_1):PSA_KEY_DERIVATION_INPUT_SECRET:"2adccada18779e7c2077ad2eb19d3f3e731385dd":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48":"":0:0:0
 
-PSA key derivation: HKDF-Expand SHA-256, RFC5869 #1 b info before secret
-depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
-derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"":0:0:0
-
 # HKDF-Expand tests: Invalid test cases
 PSA key derivation: HKDF-Expand input other secret
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
@@ -5396,6 +5388,10 @@ PSA key derivation: HKDF-Expand Invalid secret length
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SECRET:"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e500":PSA_ERROR_INVALID_ARGUMENT:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"":"":0:0:0
 
+PSA key derivation: HKDF-Expand, Info before secret
+depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
+derive_output:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SECRET:"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5":PSA_SUCCESS:0:"":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"":0:0:0
+
 # Test vectors taken from https://www.ietf.org/mail-archive/web/tls/current/msg03416.html
 PSA key derivation: TLS 1.2 PRF SHA-256, output 100+0
 depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF