diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 2c6924a..b82e67f 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -563,39 +563,43 @@
 
 PSA key policy: agreement + KDF, permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256))
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_SUCCESS
 
 PSA key policy: agreement + KDF, not permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
-agreement_key_policy:0:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256))
+agreement_key_policy:0:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ERROR_NOT_PERMITTED
 
 PSA key policy: agreement + KDF, wrong agreement algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_FFDH, PSA_ALG_HKDF(PSA_ALG_SHA_256))
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_FFDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ERROR_NOT_PERMITTED
 
 PSA key policy: agreement + KDF, wrong KDF algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_224))
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_224)):PSA_ERROR_NOT_PERMITTED
 
-PSA key policy: agreement + KDF, key only permits raw agreement
+PSA key policy: agreement + KDF, key permits raw agreement
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256))
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_SUCCESS
 
 PSA key policy: raw agreement, permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH
+raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH:PSA_SUCCESS
 
 PSA key policy: raw agreement, not permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-raw_agreement_key_policy:0:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH
+raw_agreement_key_policy:0:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH:PSA_ERROR_NOT_PERMITTED
 
 PSA key policy: raw agreement, wrong algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_FFDH
+raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_FFDH:PSA_ERROR_NOT_PERMITTED
 
-PSA key policy: raw agreement, key only permits a KDF
+PSA key policy: raw agreement, key permits raw agreement, but algorithm is not raw
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256))
+raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ERROR_NOT_SUPPORTED
+
+PSA key policy: raw agreement, key specifies KDF
+depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
+raw_agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH:PSA_ERROR_NOT_PERMITTED
 
 PSA key policy algorithm2: CTR, CBC
 depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR:MBEDTLS_CIPHER_MODE_CBC
@@ -2499,31 +2503,39 @@
 
 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_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_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_SUCCESS
+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_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_SUCCESS
+
+PSA key agreement setup: ECDH + HKDF-SHA-256: good, key algorithm broader than required
+depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_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_FAMILY_SECP_R1):PSA_ALG_ECDH:"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_SUCCESS
+
+PSA key agreement setup: ECDH + HKDF-SHA-256: key algorithm KDF mismatch
+depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_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_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_512)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_NOT_PERMITTED
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key not on curve
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_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_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ff":PSA_ERROR_INVALID_ARGUMENT
+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_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ff":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key on different curve
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_ECDH_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_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04e558dbef53eecde3d3fccfc1aea08a89a987475d12fd950d83cfa41732bc509d0d1ac43a0336def96fda41d0774a3571dcfbec7aacf3196472169e838430367f66eebe3c6e70c416dd5f0c68759dd1fff83fa40142209dff5eaad96db9e6386c":PSA_ERROR_INVALID_ARGUMENT
+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_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04e558dbef53eecde3d3fccfc1aea08a89a987475d12fd950d83cfa41732bc509d0d1ac43a0336def96fda41d0774a3571dcfbec7aacf3196472169e838430367f66eebe3c6e70c416dd5f0c68759dd1fff83fa40142209dff5eaad96db9e6386c":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key instead of private key
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_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_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
+key_agreement_setup:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH, unknown KDF
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-key_agreement_setup:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(0)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_NOT_SUPPORTED
+key_agreement_setup:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(0)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(0)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_NOT_SUPPORTED
 
 PSA key agreement setup: bad key agreement algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-key_agreement_setup:PSA_ALG_KEY_AGREEMENT(0, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
+key_agreement_setup:PSA_ALG_KEY_AGREEMENT(0, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):PSA_ALG_KEY_AGREEMENT(0, PSA_ALG_HKDF(PSA_ALG_SHA_256)):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: KDF instead of a key agreement algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-key_agreement_setup:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
+key_agreement_setup:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):PSA_ALG_HKDF(PSA_ALG_SHA_256):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":"04d12dfb5289c8d4f81208b70270398c342296970a0bccb74c736fc7554494bf6356fbf3ca366cc23e8157854c13c58d6aac23f046ada30f8353e74f33039872ab":PSA_ERROR_INVALID_ARGUMENT
 
 PSA raw key agreement: ECDH SECP256R1 (RFC 5903)
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 45916b9..fe9c1cc 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -2208,13 +2208,15 @@
                            int policy_alg,
                            int key_type_arg,
                            data_t *key_data,
-                           int exercise_alg )
+                           int exercise_alg,
+                           int expected_status_arg )
 {
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_status_t status;
+    psa_status_t expected_status = expected_status_arg;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -2228,11 +2230,7 @@
     PSA_ASSERT( psa_key_derivation_setup( &operation, exercise_alg ) );
     status = key_agreement_with_self( &operation, handle );
 
-    if( policy_alg == exercise_alg &&
-        ( policy_usage & PSA_KEY_USAGE_DERIVE ) != 0 )
-        PSA_ASSERT( status );
-    else
-        TEST_EQUAL( status, PSA_ERROR_NOT_PERMITTED );
+    TEST_EQUAL( status, expected_status );
 
 exit:
     psa_key_derivation_abort( &operation );
@@ -2283,13 +2281,15 @@
                                int policy_alg,
                                int key_type_arg,
                                data_t *key_data,
-                               int exercise_alg )
+                               int exercise_alg,
+                               int expected_status_arg )
 {
     psa_key_handle_t handle = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_status_t status;
+    psa_status_t expected_status = expected_status_arg;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -2302,11 +2302,7 @@
 
     status = raw_key_agreement_with_self( exercise_alg, handle );
 
-    if( policy_alg == exercise_alg &&
-        ( policy_usage & PSA_KEY_USAGE_DERIVE ) != 0 )
-        PSA_ASSERT( status );
-    else
-        TEST_EQUAL( status, PSA_ERROR_NOT_PERMITTED );
+    TEST_EQUAL( status, expected_status );
 
 exit:
     psa_key_derivation_abort( &operation );
@@ -5220,12 +5216,13 @@
 
 /* BEGIN_CASE */
 void key_agreement_setup( int alg_arg,
-                          int our_key_type_arg, data_t *our_key_data,
-                          data_t *peer_key_data,
+                          int our_key_type_arg, int our_key_alg_arg,
+                          data_t *our_key_data, data_t *peer_key_data,
                           int expected_status_arg )
 {
     psa_key_handle_t our_key = 0;
     psa_algorithm_t alg = alg_arg;
+    psa_algorithm_t key_alg = our_key_alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -5235,7 +5232,7 @@
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
-    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_algorithm( &attributes, key_alg );
     psa_set_key_type( &attributes, our_key_type );
     PSA_ASSERT( psa_import_key( &attributes,
                                 our_key_data->x, our_key_data->len,