diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 659205d..247c57e 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6491,74 +6491,133 @@
 PSA derive persistent key: HKDF SHA-256, exportable
 persistent_key_load_key_from_storage:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_TYPE_RAW_DATA:1024:PSA_KEY_USAGE_EXPORT:0:DERIVE_KEY
 
+PSA PAKE: uninitialized access to psa_pake_operation_t
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_SHA_256:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_UNINITIALIZED_ACCESS:PSA_ERROR_BAD_STATE
+
 PSA PAKE: invalid alg
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_SHA_256:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":PSA_ERROR_INVALID_ARGUMENT:0:0:0
+ecjpake_setup:PSA_ALG_SHA_256:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_INVALID_ARGUMENT
 
 PSA PAKE: invalid primitive type
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_DH, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":PSA_ERROR_NOT_SUPPORTED:0:0:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_DH, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive family
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_K1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":PSA_ERROR_NOT_SUPPORTED:0:0:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_K1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid primitive bits
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 128):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":PSA_ERROR_NOT_SUPPORTED:0:0:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 128):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: invalid hash
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_1:PSA_PAKE_ROLE_SERVER:0:"abcd":PSA_ERROR_NOT_SUPPORTED:0:0:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_1:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_SUPPORTED
 
-PSA PAKE: ecjpake setup server output step first
+PSA PAKE: duplicate a valid setup
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":0:0:0:0
-
-PSA PAKE: ecjpake setup server input step first
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":0:0:0:0
-
-PSA PAKE: ecjpake setup server empty password
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"":0:0:0:PSA_ERROR_BAD_STATE
-
-PSA PAKE: ecjpake setup client output step first
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":0:0:0:0
-
-PSA PAKE: ecjpake setup client input step first
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:1:"abcd":0:0:0:0
-
-PSA PAKE: ecjpake setup client empty password
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"":0:0:0:PSA_ERROR_BAD_STATE
-
-PSA PAKE: ecjpake setup client bad password key type
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_DERIVE:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":0:0:PSA_ERROR_INVALID_ARGUMENT:0
-
-PSA PAKE: ecjpake setup client bad password key usage
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_ENCRYPT:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":0:0:PSA_ERROR_NOT_PERMITTED:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_DUPLICATE_SETUP:PSA_ERROR_BAD_STATE
 
 PSA PAKE: ecjpake setup invalid role NONE
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:0:"abcd":0:PSA_ERROR_NOT_SUPPORTED:0:0
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: wrong key type password
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_HMAC:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_INVALID_ARGUMENT
+
+PSA PAKE: wrong key usage type
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:0:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_PERMITTED
+
+PSA PAKE: set invalid user
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_INVALID_USER:PSA_ERROR_INVALID_ARGUMENT
+
+PSA PAKE: set invalid peer
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_INVALID_PEER:PSA_ERROR_INVALID_ARGUMENT
+
+PSA PAKE: set user
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_SET_USER:PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: set peer
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_ERR_SET_PEER:PSA_ERROR_NOT_SUPPORTED
+
+PSA PAKE: empty server password
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"":INJECT_ERR_NONE:PSA_ERROR_BAD_STATE
+
+PSA PAKE: empty client password
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"":INJECT_ERR_NONE:PSA_ERROR_BAD_STATE
+
+PSA PAKE: invalid input
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":INJECT_EMPTY_IO_BUFFER:0
+
+PSA PAKE: unkown input step
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":INJECT_UNKNOWN_STEP:0
+
+PSA PAKE: invalid first input step
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":INJECT_INVALID_FIRST_STEP:0
+
+PSA PAKE: input buffer too large
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":INJECT_WRONG_BUFFER_SIZE:0
+
+PSA PAKE: valid input operation after a failure
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":INJECT_VALID_OPERATION_AFTER_FAILURE:0
+
+PSA PAKE: invalid output
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_EMPTY_IO_BUFFER:0
+
+PSA PAKE: unkown output step
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_UNKNOWN_STEP:0
+
+PSA PAKE: invalid first output step
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_INVALID_FIRST_STEP:0
+
+PSA PAKE: output buffer too small
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_WRONG_BUFFER_SIZE:0
+
+PSA PAKE: valid output operation after a failure
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":INJECT_VALID_OPERATION_AFTER_FAILURE:0
+
+PSA PAKE: ecjpake setup client bad password key type
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_DERIVE:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_INVALID_ARGUMENT
+
+PSA PAKE: ecjpake setup client bad password key usage
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_ENCRYPT:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":INJECT_ERR_NONE:PSA_ERROR_NOT_PERMITTED
 
 PSA PAKE: ecjpake rounds
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
-ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:0
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:INJECT_ERR_NONE
 
 PSA PAKE: ecjpake rounds, client input first
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
-ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":1:0
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":1:INJECT_ERR_NONE
 
-# This test case relies on implementation (it may need to be adjusted in the future)
-PSA PAKE: ecjpake rounds - key is destroyed after being passed to set_password_key
+PSA PAKE: ecjpake rounds, early key derivation 1
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
-ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:1
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:INJECT_ANTICIPATE_KEY_DERIVATION_1
+
+PSA PAKE: ecjpake rounds, early key derivation 2
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:INJECT_ANTICIPATE_KEY_DERIVATION_2
 
 PSA PAKE: ecjpake no input errors
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index ca1614b..8d42bf9 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1318,6 +1318,24 @@
 }
 #endif /* PSA_WANT_ALG_JPAKE */
 
+typedef enum
+{
+    INJECT_ERR_NONE = 0,
+    INJECT_ERR_UNINITIALIZED_ACCESS,
+    INJECT_ERR_DUPLICATE_SETUP,
+    INJECT_ERR_INVALID_USER,
+    INJECT_ERR_INVALID_PEER,
+    INJECT_ERR_SET_USER,
+    INJECT_ERR_SET_PEER,
+    INJECT_EMPTY_IO_BUFFER,
+    INJECT_UNKNOWN_STEP,
+    INJECT_INVALID_FIRST_STEP,
+    INJECT_WRONG_BUFFER_SIZE,
+    INJECT_VALID_OPERATION_AFTER_FAILURE,
+    INJECT_ANTICIPATE_KEY_DERIVATION_1,
+    INJECT_ANTICIPATE_KEY_DERIVATION_2,
+} ecjpake_injected_failure_t;
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -8753,11 +8771,9 @@
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
 void ecjpake_setup( int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
                     int primitive_arg, int hash_arg, int role_arg,
-                    int input_first, data_t *pw_data,
-                    int expected_status_setup_arg,
-                    int expected_status_set_role_arg,
-                    int expected_status_set_password_key_arg,
-                    int expected_status_input_output_arg)
+                    int test_input, data_t *pw_data,
+                    int inj_err_type_arg,
+                    int expected_error_arg)
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
@@ -8769,12 +8785,9 @@
     psa_pake_role_t role = role_arg;
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_status_t expected_status_setup = expected_status_setup_arg;
-    psa_status_t expected_status_set_role = expected_status_set_role_arg;
-    psa_status_t expected_status_set_password_key =
-                        expected_status_set_password_key_arg;
-    psa_status_t expected_status_input_output =
-                        expected_status_input_output_arg;
+    ecjpake_injected_failure_t inj_err_type = inj_err_type_arg;
+    psa_status_t expected_error = expected_error_arg;
+    psa_status_t status;
     unsigned char *output_buffer = NULL;
     size_t output_len = 0;
 
@@ -8799,54 +8812,90 @@
 
     PSA_ASSERT( psa_pake_abort( &operation ) );
 
-    TEST_EQUAL( psa_pake_set_user( &operation, NULL, 0 ),
-                PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_set_peer( &operation, NULL, 0 ),
-                PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_set_password_key( &operation, key ),
-                PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_set_role( &operation, role ),
-                PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
-                                 NULL, 0, NULL ),
-                PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE, NULL, 0),
-                PSA_ERROR_BAD_STATE );
-
-    PSA_ASSERT( psa_pake_abort( &operation ) );
-
-    TEST_EQUAL( psa_pake_setup( &operation, &cipher_suite ),
-                expected_status_setup );
-    if( expected_status_setup != PSA_SUCCESS )
+    if ( inj_err_type == INJECT_ERR_UNINITIALIZED_ACCESS )
+    {
+        TEST_EQUAL( psa_pake_set_user( &operation, NULL, 0 ),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
+        TEST_EQUAL( psa_pake_set_peer( &operation, NULL, 0 ),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
+        TEST_EQUAL( psa_pake_set_password_key( &operation, key ),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
+        TEST_EQUAL( psa_pake_set_role( &operation, role ),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                    NULL, 0, NULL ),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE, NULL, 0),
+                    expected_error );
+        PSA_ASSERT( psa_pake_abort( &operation ) );
         goto exit;
+    }
 
-    TEST_EQUAL( psa_pake_setup( &operation, &cipher_suite ),
-                PSA_ERROR_BAD_STATE );
-
-    TEST_EQUAL( psa_pake_set_role( &operation, role),
-                expected_status_set_role );
-    if( expected_status_set_role != PSA_SUCCESS )
+    status = psa_pake_setup( &operation, &cipher_suite );
+    if (status != PSA_SUCCESS)
+    {
+        TEST_EQUAL( status, expected_error );
         goto exit;
+    }
+
+    if( inj_err_type == INJECT_ERR_DUPLICATE_SETUP )
+    {
+        TEST_EQUAL( psa_pake_setup( &operation, &cipher_suite ),
+                expected_error );
+        goto exit;
+    }
+
+    status = psa_pake_set_role( &operation, role);
+    if ( status != PSA_SUCCESS )
+    {
+        TEST_EQUAL( status, expected_error );
+        goto exit;
+    }
 
     if( pw_data->len > 0 )
     {
-        TEST_EQUAL( psa_pake_set_password_key( &operation, key ),
-                    expected_status_set_password_key );
-        if( expected_status_set_password_key != PSA_SUCCESS )
+        status = psa_pake_set_password_key( &operation, key );
+        if ( status != PSA_SUCCESS )
+        {
+            TEST_EQUAL( status, expected_error );
             goto exit;
+        }
     }
 
-    TEST_EQUAL( psa_pake_set_user( &operation, NULL, 0 ),
-                PSA_ERROR_INVALID_ARGUMENT );
-    TEST_EQUAL( psa_pake_set_peer( &operation, NULL, 0 ),
-                PSA_ERROR_INVALID_ARGUMENT );
+    if ( inj_err_type == INJECT_ERR_INVALID_USER )
+    {
+        TEST_EQUAL( psa_pake_set_user( &operation, NULL, 0 ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        goto exit;
+    }
 
-    const uint8_t unsupported_id[] = "abcd";
+    if ( inj_err_type == INJECT_ERR_INVALID_PEER )
+    {
+        TEST_EQUAL( psa_pake_set_peer( &operation, NULL, 0 ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        goto exit;
+    }
 
-    TEST_EQUAL( psa_pake_set_user( &operation, unsupported_id, 4 ),
-                PSA_ERROR_NOT_SUPPORTED );
-    TEST_EQUAL( psa_pake_set_peer( &operation, unsupported_id, 4 ),
-                PSA_ERROR_NOT_SUPPORTED );
+    if ( inj_err_type == INJECT_ERR_SET_USER )
+    {
+        const uint8_t unsupported_id[] = "abcd";
+        TEST_EQUAL( psa_pake_set_user( &operation, unsupported_id, 4 ),
+                    PSA_ERROR_NOT_SUPPORTED );
+        goto exit;
+    }
+
+    if ( inj_err_type == INJECT_ERR_SET_PEER )
+    {
+        const uint8_t unsupported_id[] = "abcd";
+        TEST_EQUAL( psa_pake_set_peer( &operation, unsupported_id, 4 ),
+                    PSA_ERROR_NOT_SUPPORTED );
+        goto exit;
+    }
 
     const size_t size_key_share = PSA_PAKE_INPUT_SIZE( alg, primitive,
                                                 PSA_PAKE_STEP_KEY_SHARE );
@@ -8855,85 +8904,109 @@
     const size_t size_zk_proof = PSA_PAKE_INPUT_SIZE( alg, primitive,
                                                 PSA_PAKE_STEP_ZK_PROOF );
 
-    /* First round */
-    if( input_first )
+    if ( test_input )
     {
-        /* Invalid parameters (input) */
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
-                                    NULL, 0 ),
-                    PSA_ERROR_INVALID_ARGUMENT );
-        /* Invalid parameters (step) */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                 key, pw_data->len ) , 0 );
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
-                                    output_buffer, size_zk_proof ),
-                    PSA_ERROR_INVALID_ARGUMENT );
-        /* Invalid first step */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                 key, pw_data->len ), 0 );
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
-                                    output_buffer, size_zk_proof ),
-                    PSA_ERROR_BAD_STATE );
-
-        /* Possibly valid */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                 key, pw_data->len ), 0 );
-        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
-                                    output_buffer, size_key_share ),
-                    expected_status_input_output);
-
-        if( expected_status_input_output == PSA_SUCCESS )
+        if ( inj_err_type == INJECT_EMPTY_IO_BUFFER )
         {
-            /* Buffer too large */
-            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
-                                    output_buffer, size_zk_public + 1 ),
-                        PSA_ERROR_INVALID_ARGUMENT );
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF, NULL, 0 ),
+                                        PSA_ERROR_INVALID_ARGUMENT );
+            goto exit;
+        }
 
-            /* The operation's state should be invalidated at this point */
+        if ( inj_err_type == INJECT_UNKNOWN_STEP )
+        {
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+                                        output_buffer, size_zk_proof ),
+                                        PSA_ERROR_INVALID_ARGUMENT );
+            goto exit;
+        }
+
+        if ( inj_err_type == INJECT_INVALID_FIRST_STEP )
+        {
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                        output_buffer, size_zk_proof ),
+                                        PSA_ERROR_BAD_STATE );
+            goto exit;
+        }
+
+        status = psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                                        output_buffer, size_key_share );
+        if ( status != PSA_SUCCESS )
+        {
+            TEST_EQUAL( status, expected_error);
+            goto exit;
+        }
+
+        if ( inj_err_type == INJECT_WRONG_BUFFER_SIZE )
+        {
+            TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, size_zk_public + 1 ),
+                                        PSA_ERROR_INVALID_ARGUMENT );
+            goto exit;
+        }
+
+        if ( inj_err_type == INJECT_VALID_OPERATION_AFTER_FAILURE )
+        {
+            // Just trigger any kind of error. We don't care about the result here
+            psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, size_zk_public + 1 );
             TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
                                         output_buffer, size_zk_public ),
-                        PSA_ERROR_BAD_STATE );
+                                        PSA_ERROR_BAD_STATE );
+            goto exit;
         }
-    }
-    else
-    {
-        /* Invalid parameters (output) */
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
-                                     NULL, 0, NULL ),
-                    PSA_ERROR_INVALID_ARGUMENT );
-        /* Invalid parameters (step) */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                key, pw_data->len ), 0 );
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
-                                     output_buffer, buf_size, &output_len ),
-                    PSA_ERROR_INVALID_ARGUMENT );
-        /* Invalid first step */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                 key, pw_data->len ), 0 );
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
-                                     output_buffer, buf_size, &output_len ),
-                    PSA_ERROR_BAD_STATE );
-
-        /* Possibly valid */
-        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
-                                 key, pw_data->len ), 0 );
-        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
-                                     output_buffer, buf_size, &output_len ),
-                    expected_status_input_output );
-
-        if( expected_status_input_output == PSA_SUCCESS )
+    } else {
+        if ( inj_err_type == INJECT_EMPTY_IO_BUFFER )
         {
-            TEST_ASSERT( output_len > 0 );
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                        NULL, 0, NULL ),
+                                        PSA_ERROR_INVALID_ARGUMENT );
+            goto exit;
+        }
 
-            /* Buffer too small */
-            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
-                                         output_buffer, size_zk_public - 1, &output_len ),
-                        PSA_ERROR_BUFFER_TOO_SMALL );
+        if ( inj_err_type == INJECT_UNKNOWN_STEP )
+        {
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
+                                        output_buffer, buf_size, &output_len ),
+                                        PSA_ERROR_INVALID_ARGUMENT );
+            goto exit;
+        }
 
-            /* The operation's state should be invalidated at this point */
+        if ( inj_err_type == INJECT_INVALID_FIRST_STEP )
+        {
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
+                                        output_buffer, buf_size, &output_len ),
+                                        PSA_ERROR_BAD_STATE );
+            goto exit;
+        }
+
+        status = psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
+                                    output_buffer, buf_size, &output_len );
+        if ( status != PSA_SUCCESS )
+        {
+            TEST_EQUAL( status, expected_error);
+            goto exit;
+        }
+
+        TEST_ASSERT( output_len > 0 );
+
+        if ( inj_err_type == INJECT_WRONG_BUFFER_SIZE )
+        {
             TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
-                                         output_buffer, buf_size, &output_len ),
-                        PSA_ERROR_BAD_STATE );
+                                        output_buffer, size_zk_public - 1, &output_len ),
+                                        PSA_ERROR_BUFFER_TOO_SMALL );
+            goto exit;
+        }
+
+        if ( inj_err_type == INJECT_VALID_OPERATION_AFTER_FAILURE )
+        {
+            // Just trigger any kind of error. We don't care about the result here
+            psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, size_zk_public - 1, &output_len );
+            TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
+                                        output_buffer, buf_size, &output_len ),
+                                        PSA_ERROR_BAD_STATE );
+            goto exit;
         }
     }
 
@@ -9000,7 +9073,8 @@
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
 void ecjpake_rounds( int alg_arg, int primitive_arg, int hash_arg,
                      int derive_alg_arg, data_t *pw_data,
-                     int client_input_first, int destroy_key )
+                     int client_input_first, int destroy_key,
+                     int client_input_first, int inj_err_type_arg )
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t server = psa_pake_operation_init();
@@ -9014,6 +9088,7 @@
                             PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_derivation_operation_t client_derive =
                             PSA_KEY_DERIVATION_OPERATION_INIT;
+    ecjpake_injected_failure_t inj_err_type = inj_err_type_arg;
 
     PSA_INIT( );
 
@@ -9054,19 +9129,27 @@
     if( destroy_key == 1 )
         psa_destroy_key( key );
 
-    TEST_EQUAL( psa_pake_get_implicit_key( &server, &server_derive ),
+    if( inj_err_type == INJECT_ANTICIPATE_KEY_DERIVATION_1 )
+    {
+        TEST_EQUAL( psa_pake_get_implicit_key( &server, &server_derive ),
                 PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_get_implicit_key( &client, &client_derive ),
+        TEST_EQUAL( psa_pake_get_implicit_key( &client, &client_derive ),
                 PSA_ERROR_BAD_STATE );
+        goto exit;
+    }
 
     /* First round */
     ecjpake_do_round( alg, primitive_arg, &server, &client,
                       client_input_first, 1, 0 );
 
-    TEST_EQUAL( psa_pake_get_implicit_key( &server, &server_derive ),
+    if ( inj_err_type == INJECT_ANTICIPATE_KEY_DERIVATION_2 )
+    {
+        TEST_EQUAL( psa_pake_get_implicit_key( &server, &server_derive ),
                 PSA_ERROR_BAD_STATE );
-    TEST_EQUAL( psa_pake_get_implicit_key( &client, &client_derive ),
+        TEST_EQUAL( psa_pake_get_implicit_key( &client, &client_derive ),
                 PSA_ERROR_BAD_STATE );
+        goto exit;
+    }
 
     /* Second round */
     ecjpake_do_round( alg, primitive_arg, &server, &client,