diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 55ab3d2245..ad958ad7fb 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -7166,39 +7166,39 @@ derive_key:PSA_ALG_PBKDF2_AES_CMAC_PRF_128:"706173737764":"01":"73616c74":PSA_KE
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: good
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: good, key algorithm broader than required
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: key algorithm KDF mismatch
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_SHA_512:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS:PSA_ERROR_NOT_PERMITTED
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key not on curve
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key on different curve
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ECC_SECP_R1_384
-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
+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_SUCCESS:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: public key instead of private key
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: ECDH, unknown KDF
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:PSA_WANT_ECC_SECP_R1_256
-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
+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:0
 
 PSA key agreement setup: bad key agreement algorithm
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key agreement setup: KDF instead of a key agreement algorithm
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
-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
+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_SUCCESS:PSA_ERROR_INVALID_ARGUMENT
 
 PSA (raw) key agreement: ECDH SECP256R1 (RFC 5903)
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_DERIVE:PSA_WANT_ECC_SECP_R1_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b416311bd7..0f679f0581 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -9797,7 +9797,8 @@ exit:
 void key_agreement_setup(int alg_arg,
                          int our_key_type_arg, int our_key_alg_arg,
                          data_t *our_key_data, data_t *peer_key_data,
-                         int expected_status_arg)
+                         int expected_status_setup_arg,
+                         int expected_status_keyag_arg)
 {
     mbedtls_svc_key_id_t our_key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_algorithm_t alg = alg_arg;
@@ -9805,8 +9806,8 @@ void key_agreement_setup(int alg_arg,
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_key_derivation_operation_t operation = psa_key_derivation_operation_init_short();
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_status_t expected_status = expected_status_arg;
-    psa_status_t status;
+    psa_status_t expected_status_setup = expected_status_setup_arg;
+    psa_status_t expected_status_keyag = expected_status_keyag_arg;
 
     PSA_ASSERT(psa_crypto_init());
 
@@ -9817,21 +9818,18 @@ void key_agreement_setup(int alg_arg,
                               our_key_data->x, our_key_data->len,
                               &our_key));
 
-    /* The tests currently include inputs that should fail at either step.
-     * Test cases that fail at the setup step should be changed to call
-     * key_derivation_setup instead, and this function should be renamed
-     * to key_agreement_fail. */
-    status = psa_key_derivation_setup(&operation, alg);
-    if (status == PSA_SUCCESS) {
-        TEST_EQUAL(psa_key_derivation_key_agreement(
-                       &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
-                       our_key,
-                       peer_key_data->x, peer_key_data->len),
-                   expected_status);
-    } else {
-        TEST_ASSERT(status == expected_status);
+    TEST_EQUAL(psa_key_derivation_setup(&operation, alg),
+               expected_status_setup);
+    if (expected_status_setup != PSA_SUCCESS) {
+        goto exit;
     }
 
+    TEST_EQUAL(psa_key_derivation_key_agreement(
+                   &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
+                   our_key,
+                   peer_key_data->x, peer_key_data->len),
+               expected_status_keyag);
+
 exit:
     psa_key_derivation_abort(&operation);
     psa_destroy_key(our_key);