diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index a32da21..6522fe5 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -3050,8 +3050,6 @@
             /* --- psa_pake_input (driver: setup, input) --- */
             mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
-                   sizeof(mbedtls_test_driver_pake_hooks.hits));
             TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
                                       input_buffer, size_key_share),
                        expected_status);
@@ -3063,8 +3061,6 @@
             /* --- psa_pake_output (driver: setup, output) --- */
             mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
-                   sizeof(mbedtls_test_driver_pake_hooks.hits));
             TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
                                        output_buffer, output_size, &output_len),
                        expected_status);
@@ -3073,11 +3069,9 @@
             break;
 
         case 2: /* input */
-            /* --- psa_pake_input (driver: setup, input, (abort)) --- */
+            /* --- psa_pake_input (driver: setup, input, abort) --- */
             mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
-                   sizeof(mbedtls_test_driver_pake_hooks.hits));
             TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
                                       input_buffer, size_key_share),
                        expected_status);
@@ -3091,8 +3085,6 @@
             /* --- psa_pake_output (driver: setup, output, (abort)) --- */
             mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-            memset(&mbedtls_test_driver_pake_hooks.hits, 0,
-                   sizeof(mbedtls_test_driver_pake_hooks.hits));
             if (forced_output->len > 0) {
                 mbedtls_test_driver_pake_hooks.forced_output = forced_output->x;
                 mbedtls_test_driver_pake_hooks.forced_output_length = forced_output->len;
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index 3be249f..1a25a8c 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -194,14 +194,29 @@
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256
 ecjpake_size_macros:
 
-PSA PAKE: input getters: ok #1
-pake_input_getters:"aabbccddee":PSA_PAKE_ROLE_SERVER:5:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+PSA PAKE: input getters password: ok #1
+pake_input_getters_password:"aabbccddee":5:PSA_SUCCESS:PSA_SUCCESS
 
-PSA PAKE: input getters: ok #2
-pake_input_getters:"ddccbbaa":PSA_PAKE_ROLE_CLIENT:5:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_512:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+PSA PAKE: input getters password: ok #2
+pake_input_getters_password:"11223344556677889900":10:PSA_SUCCESS:PSA_SUCCESS
 
-PSA PAKE: input getters: buffer for password to small
-pake_input_getters:"aabbccddee":PSA_PAKE_ROLE_SERVER:4:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ERROR_BUFFER_TOO_SMALL:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+PSA PAKE: input getters password: buffer to small
+pake_input_getters_password:"aabbccddee":4:PSA_ERROR_BUFFER_TOO_SMALL:PSA_SUCCESS
 
-PSA PAKE: input getters: inputs not ready
-pake_input_getters:"":0:5:0:0:0:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+PSA PAKE: input getters password: not set
+pake_input_getters_password:"":0:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
+
+PSA PAKE: input getters cipher suite: ok
+pake_input_getters_cipher_suite:PSA_SUCCESS:1
+
+PSA PAKE: input getters cipher suite: not set
+pake_input_getters_cipher_suite:PSA_ERROR_BAD_STATE:0
+
+PSA PAKE: input getters role client: ok
+pake_input_getters_role:PSA_PAKE_ROLE_CLIENT:PSA_SUCCESS
+
+PSA PAKE: input getters role server: ok
+pake_input_getters_role:PSA_PAKE_ROLE_SERVER:PSA_SUCCESS
+
+PSA PAKE: input getters role: not set
+pake_input_getters_role:PSA_PAKE_ROLE_NONE:PSA_ERROR_BAD_STATE
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index 1c3b328..3bb441f 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -909,71 +909,54 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_JPAKE */
-void pake_input_getters(data_t *password, int role_arg, int password_buffer_size,
-                        int alg_arg, int primitive_arg, int hash_arg,
-                        int expected_status_pass,  int expected_status_pass_len,
-                        int expected_status_role,  int expected_status_cs)
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_password(data_t *password, int password_buffer_size,
+                                 int expected_status_pass, int expected_status_pass_len)
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
-    psa_pake_role_t role = role_arg;
-    psa_algorithm_t alg = alg_arg;
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_pake_role_t role_ret = PSA_PAKE_ROLE_NONE;
     uint8_t password_ret[20] = { 0 }; // max key length is 20 bytes
     size_t password_len_ret = 0;
-    psa_pake_cipher_suite_t cipher_suite_ret = psa_pake_cipher_suite_init();
     size_t buffer_len_ret = 0;
 
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
     PSA_INIT();
 
-    /* alg equal to 0 indicates case when inputs are not set yet. */
-    if (alg != 0) {
-        psa_pake_cs_set_algorithm(&cipher_suite, alg);
-        psa_pake_cs_set_primitive(&cipher_suite, primitive_arg);
-        psa_pake_cs_set_hash(&cipher_suite, hash_arg);
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
 
-        psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
-        psa_set_key_algorithm(&attributes, alg);
-        psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
+    psa_set_key_algorithm(&attributes, PSA_ALG_JPAKE);
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
 
+    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+
+    if (password_buffer_size > 0) {
         PSA_ASSERT(psa_import_key(&attributes, password->x, password->len, &key));
-
-        PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
         PSA_ASSERT(psa_pake_set_password_key(&operation, key));
-        PSA_ASSERT(psa_pake_set_role(&operation, role));
     }
 
     TEST_EQUAL(psa_crypto_driver_pake_get_password_len(&operation.data.inputs, &password_len_ret),
                expected_status_pass_len);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
-                                                   (uint8_t *) &password_ret,
-                                                   password_buffer_size, &buffer_len_ret),
-               expected_status_pass);
-
-    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
-               expected_status_role);
-
-    TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
-               expected_status_cs);
-
     if (expected_status_pass_len == PSA_SUCCESS) {
         TEST_EQUAL(password_len_ret, password->len);
-    }
 
-    if (expected_status_pass == PSA_SUCCESS) {
-        PSA_ASSERT(memcmp(password_ret, password->x, password->len));
-    }
+        TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
+                                                       (uint8_t *) &password_ret,
+                                                       password_buffer_size, &buffer_len_ret),
+                   expected_status_pass);
 
-    if (expected_status_role == PSA_SUCCESS) {
-        TEST_EQUAL(role_ret, role);
-    }
-
-    if (expected_status_pass == PSA_SUCCESS) {
-        PSA_ASSERT(memcmp(&cipher_suite_ret, &cipher_suite, sizeof(cipher_suite)));
+        if (expected_status_pass == PSA_SUCCESS) {
+            TEST_EQUAL(buffer_len_ret, password->len);
+            PSA_ASSERT(memcmp(password_ret, password->x, buffer_len_ret));
+        }
     }
 
 exit:
@@ -982,3 +965,74 @@
     PSA_DONE();
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_cipher_suite(int expected_status, int setup_done)
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_cipher_suite_t cipher_suite_ret = psa_pake_cipher_suite_init();
+
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
+    PSA_INIT();
+
+    if (setup_done == 1) {
+        psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+        psa_pake_cs_set_primitive(&cipher_suite, primitive);
+        psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+        PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+    }
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
+               expected_status);
+
+    if (expected_status == PSA_SUCCESS) {
+        PSA_ASSERT(memcmp(&cipher_suite_ret, &cipher_suite, sizeof(cipher_suite)));
+    }
+
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
+void pake_input_getters_role(int role_arg, int expected_status)
+{
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_pake_operation_t operation = psa_pake_operation_init();
+    psa_pake_role_t role_ret = PSA_PAKE_ROLE_NONE;
+    psa_pake_role_t role = role_arg;
+
+    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
+        PSA_PAKE_PRIMITIVE_TYPE_ECC,
+        PSA_ECC_FAMILY_SECP_R1, 256);
+
+    PSA_INIT();
+
+    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
+    psa_pake_cs_set_primitive(&cipher_suite, primitive);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
+
+    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+
+    if (role != PSA_PAKE_ROLE_NONE) {
+        PSA_ASSERT(psa_pake_set_role(&operation, role));
+    }
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
+               expected_status);
+
+    if (expected_status == PSA_SUCCESS) {
+        TEST_EQUAL(role_ret, role);
+    }
+
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */