diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 4a052c4..fd33eba 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5799,36 +5799,44 @@
 
 PSA key derivation: ECJPAKE to PMS, no input
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, bad input format
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 #NIST CAVS 11.0 SHA-256 ShortMSG vector for L=256
 PSA key derivation: ECJPAKE to PMS, good case
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
 
 PSA key derivation: ECJPAKE to PMS, bad derivation step
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SEED:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SEED:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, capacity 1 byte too big
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:33:PSA_ERROR_INVALID_ARGUMENT:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, capacity 1 byte too small
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:31:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_ERROR_INSUFFICIENT_DATA
 
 PSA key derivation: ECJPAKE to PMS, output too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, output too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INSUFFICIENT_DATA
 
 PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
 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 0162c07..355cba5 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7235,7 +7235,9 @@
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS:MBEDTLS_SHA256_C */
 void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
-                            int derivation_step, data_t *expected_output,
+                            int derivation_step,
+                            int capacity, int expected_capacity_status,
+                            data_t *expected_output,
                             int expected_output_status )
 {
     psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
@@ -7248,8 +7250,8 @@
     PSA_ASSERT( psa_crypto_init() );
 
     PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-    PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
-                                                 expected_output->len ) );
+    TEST_EQUAL( psa_key_derivation_set_capacity( &operation, capacity ),
+        (psa_status_t) expected_capacity_status );
 
     TEST_EQUAL( psa_key_derivation_input_bytes( &operation,
                 step, input->x, input->len ),