diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 38bc64c..c8b229c 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6448,51 +6448,59 @@
 
 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_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
+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
 
 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_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
+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
 
 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_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
+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
 
 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_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
+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
 
 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_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
+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
 
 PSA PAKE: ecjpake setup server 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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"abcd":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":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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:1:"abcd":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: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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_SERVER:0:"":PSA_ERROR_BAD_STATE
+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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"abcd":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_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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:1:"abcd":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_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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_CLIENT:0:"":PSA_ERROR_BAD_STATE
+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
 
 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_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_PAKE_ROLE_NONE:0:"abcd":PSA_ERROR_NOT_SUPPORTED
+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
 
 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
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 4ca308d..c74acf6 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8657,19 +8657,29 @@
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
-void ecjpake_setup( int alg_arg, int primitive_arg, int hash_arg, int role_arg,
+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_arg )
+                    int expected_status_setup_arg,
+                    int expected_status_set_role_arg,
+                    int expected_status_set_password_key_arg,
+                    int expected_status_input_output_arg)
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
     psa_algorithm_t alg = alg_arg;
+    psa_key_type_t key_type_pw = key_type_pw_arg;
+    psa_key_usage_t key_usage_pw = key_usage_pw_arg;
     psa_algorithm_t hash_alg = hash_arg;
     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 = expected_status_arg;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    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;
     unsigned char *output_buffer = NULL;
     size_t output_len = 0;
 
@@ -8681,9 +8691,9 @@
 
     if( pw_data->len > 0 )
     {
-        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+        psa_set_key_usage_flags( &attributes, key_usage_pw );
         psa_set_key_algorithm( &attributes, alg );
-        psa_set_key_type( &attributes, PSA_KEY_TYPE_PASSWORD );
+        psa_set_key_type( &attributes, key_type_pw );
         PSA_ASSERT( psa_import_key( &attributes, pw_data->x, pw_data->len,
                                     &key ) );
     }
@@ -8710,37 +8720,25 @@
 
     PSA_ASSERT( psa_pake_abort( &operation ) );
 
-    status = psa_pake_setup( &operation, &cipher_suite );
-    if( status != PSA_SUCCESS )
-    {
-        TEST_EQUAL( status, expected_status );
+    TEST_EQUAL( psa_pake_setup( &operation, &cipher_suite ),
+                expected_status_setup );
+    if( expected_status_setup != PSA_SUCCESS )
         goto exit;
-    }
-    else
-        PSA_ASSERT( status );
 
     TEST_EQUAL( psa_pake_setup( &operation, &cipher_suite ),
                 PSA_ERROR_BAD_STATE );
 
-    status = psa_pake_set_role( &operation, role );
-    if( status != PSA_SUCCESS )
-    {
-        TEST_EQUAL( status, expected_status );
+    TEST_EQUAL( psa_pake_set_role( &operation, role),
+                expected_status_set_role );
+    if( expected_status_set_role != PSA_SUCCESS )
         goto exit;
-    }
-    else
-        PSA_ASSERT( status );
 
     if( pw_data->len > 0 )
     {
-        status = psa_pake_set_password_key( &operation, key );
-        if( status != PSA_SUCCESS )
-        {
-            TEST_EQUAL( status, expected_status );
+        TEST_EQUAL( psa_pake_set_password_key( &operation, key ),
+                    expected_status_set_password_key );
+        if( expected_status_set_password_key != PSA_SUCCESS )
             goto exit;
-        }
-        else
-            PSA_ASSERT( status );
     }
 
     TEST_EQUAL( psa_pake_set_user( &operation, NULL, 0 ),
@@ -8772,9 +8770,9 @@
 
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                     output_buffer, 66 ),
-                    expected_status);
+                    expected_status_input_output);
 
-        if( expected_status == PSA_SUCCESS )
+        if( expected_status_input_output == PSA_SUCCESS )
         {
             /* Buffer too large */
             TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PUBLIC,
@@ -8803,9 +8801,9 @@
 
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                      output_buffer, 512, &output_len ),
-                    expected_status );
+                    expected_status_input_output );
 
-        if( expected_status == PSA_SUCCESS )
+        if( expected_status_input_output == PSA_SUCCESS )
         {
             TEST_ASSERT( output_len > 0 );