diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index bcdf68336..76af6b66d 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7478,16 +7478,20 @@ psa_status_t psa_pake_set_role(
         goto exit;
     }
 
-    if (role != PSA_PAKE_ROLE_NONE &&
-        role != PSA_PAKE_ROLE_FIRST &&
-        role != PSA_PAKE_ROLE_SECOND &&
-        role != PSA_PAKE_ROLE_CLIENT &&
-        role != PSA_PAKE_ROLE_SERVER) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto exit;
+    switch (operation->alg) {
+#if defined(PSA_WANT_ALG_JPAKE)
+        case PSA_ALG_JPAKE:
+            if (role == PSA_PAKE_ROLE_NONE) {
+                return PSA_SUCCESS;
+            }
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            break;
+#endif
+        default:
+            (void) role;
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto exit;
     }
-
-    status = PSA_ERROR_NOT_SUPPORTED;
 exit:
     psa_pake_abort(operation);
     return status;
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index eeef53a27..62157037d 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -28,7 +28,7 @@ ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_
 
 PSA PAKE: ecjpake setup role
 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:"client":"server":0:ERR_INJECT_SET_ROLE: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:"client":"server":0:ERR_INJECT_SET_ROLE:PSA_ERROR_INVALID_ARGUMENT
 
 PSA PAKE: wrong password key type
 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_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index 07d4e9f1f..ebfe843c6 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -42,6 +42,7 @@ typedef enum {
     ERR_IN_SETUP,
     ERR_IN_SET_USER,
     ERR_IN_SET_PEER,
+    ERR_IN_SET_ROLE,
     ERR_IN_SET_PASSWORD_KEY,
     ERR_IN_INPUT,
     ERR_IN_OUTPUT,
@@ -614,6 +615,9 @@ void ecjpake_setup(int alg_arg, int key_type_pw_arg, int key_usage_pw_arg,
     SETUP_CONDITIONAL_CHECK_STEP(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER),
                                  ERR_INJECT_SET_ROLE);
 
+    SETUP_ALWAYS_CHECK_STEP(psa_pake_set_role(&operation, PSA_PAKE_ROLE_NONE),
+                                 ERR_IN_SET_ROLE);
+
     SETUP_ALWAYS_CHECK_STEP(psa_pake_set_user(&operation, user, user_len),
                                  ERR_IN_SET_USER);