diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index 73c569d..fa7aa7b 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -821,38 +821,54 @@
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
 aead_decrypt_setup:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c87281":"4365847fe0b7b7fbed325953df344a96":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY
 
-PSA PAKE setup transparent driver: in-driver success
+PSA PAKE transparent driver: setup(via input) in-driver forced status
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-pake_setup:"abcd":PSA_SUCCESS:PSA_SUCCESS
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:0
 
-PSA PAKE setup transparent driver: in-driver forced error
+PSA PAKE transparent driver: setup(via output) in-driver forced status
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-pake_setup:"abcd":PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:1
 
-PSA PAKE setup transparent driver: fallback
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:MBEDTLS_PSA_BUILTIN_PAKE
-pake_setup:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+PSA PAKE transparent driver: input in-driver forced status
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:2
 
-PSA PAKE setup transparent driver: fallback not available
+PSA PAKE transparent driver: output in-driver forced status
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:3
+
+PSA PAKE transparent driver: output in-driver forced output
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+pake_operations:"abcd":PSA_SUCCESS:PSA_SUCCESS:"1234":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:3
+
+PSA PAKE transparent driver: get_key in-driver forced status
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:4
+
+PSA PAKE transparent driver: abort in-driver forced status
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:5
+
+PSA PAKE transparent driver: setup(via input) fallback not available
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:!MBEDTLS_PSA_BUILTIN_PAKE
-pake_setup:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_ERROR_NOT_SUPPORTED
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:0
 
-PSA PAKE operations transparent driver: in-driver success
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-pake_operations:"abcd":PSA_SUCCESS:"":PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_BAD_STATE
+PSA PAKE transparent driver: setup(via output) fallback not available
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:1
 
-PSA PAKE operations transparent driver: in-driver forced status
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-pake_operations:"abcd":PSA_SUCCESS:"1234":PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY
+PSA PAKE transparent driver: input fallback not available
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:2
 
-PSA PAKE operations transparent driver: fallback
-depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:MBEDTLS_PSA_BUILTIN_PAKE
-pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:"":PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_BAD_STATE
+PSA PAKE transparent driver: output fallback not available
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_SUCCESS:PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:3
 
-PSA PAKE: ecjpake rounds transparent driver: in-driver
+PSA PAKE: ecjpake rounds transparent driver: in-driver success
 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
 
-PSA PAKE: ecjpake rounds transparent driver: fallback
+PSA PAKE: ecjpake rounds transparent driver: fallback success
 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:MBEDTLS_PSA_BUILTIN_PAKE
 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
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index cfbcccb..2e1c626 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -3,8 +3,8 @@
 
 /* Auxiliary variables for pake tests.
    Global to silent the compiler when unused. */
-size_t pake_expected_hit_count;
-int pake_in_driver;
+size_t pake_expected_hit_count = 0;
+int pake_in_driver = 0;
 
 #if defined(PSA_WANT_ALG_JPAKE)
 static void ecjpake_do_round(psa_algorithm_t alg, unsigned int primitive,
@@ -142,6 +142,9 @@
                 TEST_EQUAL(status, PSA_SUCCESS);
             }
 
+            /* Adjust for indirect client driver setup in first pake_output call. */
+            pake_expected_hit_count++;
+
             /* Client first round Output */
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer1 + buffer1_off,
@@ -2974,81 +2977,27 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void pake_setup(data_t *pw_data, int forced_status_arg, int expected_status_arg)
-{
-    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
-    psa_status_t forced_status = forced_status_arg;
-    psa_status_t expected_status = expected_status_arg;
-    psa_pake_operation_t operation = psa_pake_operation_init();
-    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
-    psa_key_type_t key_type_pw = PSA_KEY_TYPE_PASSWORD;
-    psa_key_usage_t key_usage_pw = PSA_KEY_USAGE_DERIVE;
-    psa_algorithm_t alg = PSA_ALG_JPAKE;
-    psa_algorithm_t hash_alg = PSA_ALG_SHA_256;
-    psa_pake_primitive_t primitive_arg = PSA_PAKE_PRIMITIVE(
-        PSA_PAKE_PRIMITIVE_TYPE_ECC,
-        PSA_ECC_FAMILY_SECP_R1, 256);
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_init();
-
-    PSA_INIT();
-
-    if (pw_data->len > 0) {
-        psa_set_key_usage_flags(&attributes, key_usage_pw);
-        psa_set_key_algorithm(&attributes, alg);
-        psa_set_key_type(&attributes, key_type_pw);
-        PSA_ASSERT(psa_import_key(&attributes, pw_data->x, pw_data->len,
-                                  &key));
-    }
-
-    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_alg);
-
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-
-    TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
-               expected_status);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes(&attributes);
-    psa_destroy_key(key);
-    mbedtls_test_driver_pake_hooks =
-        mbedtls_test_driver_pake_hooks_init();
-    PSA_DONE();
-}
-/* END_CASE */
-
-/* BEGIN_CASE */
-void pake_operations(data_t *pw_data, int forced_status_setup_arg, data_t *forced_output,
-                     int forced_status_arg, int expected_status_set_user_arg,
-                     int expected_status_set_role_arg, int expected_status_set_peer_arg,
-                     int expected_status_set_password_arg, int expected_status_input_arg,
-                     int expected_status_abort_arg, int expected_status_output_arg,
-                     int expected_status_get_key_arg)
+void pake_operations(data_t *pw_data, int forced_status_setup_arg, int forced_status_arg,
+                     data_t *forced_output, int expected_status_setup_arg,
+                     int expected_status_input_arg, int expected_status_output_arg,
+                     int expected_status_get_key_arg, int expected_status_abort_arg,
+                     int fut)
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t forced_status = forced_status_arg;
     psa_status_t forced_status_setup = forced_status_setup_arg;
-    psa_status_t expected_status_set_user = expected_status_set_user_arg;
-    psa_status_t expected_status_set_role = expected_status_set_role_arg;
-    psa_status_t expected_status_set_peer = expected_status_set_peer_arg;
-    psa_status_t expected_status_set_password = expected_status_set_password_arg;
+    psa_status_t expected_status_setup = expected_status_setup_arg;
     psa_status_t expected_status_input = expected_status_input_arg;
-    psa_status_t expected_status_abort = expected_status_abort_arg;
     psa_status_t expected_status_output = expected_status_output_arg;
     psa_status_t expected_status_get_key = expected_status_get_key_arg;
+    psa_status_t expected_status_abort = expected_status_abort_arg;
     psa_pake_operation_t operation = psa_pake_operation_init();
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_key_type_t key_type_pw = PSA_KEY_TYPE_PASSWORD;
     psa_key_usage_t key_usage_pw = PSA_KEY_USAGE_DERIVE;
     psa_algorithm_t alg = PSA_ALG_JPAKE;
     psa_algorithm_t hash_alg = PSA_ALG_SHA_256;
+    int in_driver = 1;
     psa_key_derivation_operation_t implicit_key =
         PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
@@ -3056,9 +3005,6 @@
         PSA_ECC_FAMILY_SECP_R1, 256);
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_init();
-    const unsigned char *user = (const unsigned char *) "user";
-    const uint8_t peer[] = "abcd";
-    uint32_t expected_hits = 1;
     unsigned char *input_buffer = NULL;
     const size_t size_key_share = PSA_PAKE_INPUT_SIZE(alg, primitive,
                                                       PSA_PAKE_STEP_KEY_SHARE);
@@ -3077,9 +3023,8 @@
                                      PSA_PAKE_STEP_KEY_SHARE));
     memset(output_buffer, 0x55, output_size);
 
-    /* Transparent driver is not available (fallback). */
-    if (forced_status_setup == PSA_ERROR_NOT_SUPPORTED) {
-        expected_hits = 0;
+    if (forced_status_setup_arg == PSA_ERROR_NOT_SUPPORTED) {
+        in_driver = 0;
     }
 
     PSA_INIT();
@@ -3097,99 +3042,9 @@
     psa_pake_cs_set_hash(&cipher_suite, hash_alg);
 
     mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
-    TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
-               PSA_SUCCESS);
 
-    /* --- psa_pake_set_user --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
+    /* Collecting input stage (no driver entry points) */
 
-    TEST_EQUAL(psa_pake_set_user(&operation, user, 4),
-               expected_status_set_user);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* psa_pake_set_user is unsupported (after this call operation is aborted)
-       we need to reinitialize object. */
-    if (mbedtls_test_driver_pake_hooks.forced_status == PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
-        TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
-                   PSA_SUCCESS);
-    }
-
-    /* --- psa_pake_set_peer --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_set_peer(&operation, peer, 4),
-               expected_status_set_peer);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* psa_pake_set_user is unsupported (after this call operation is aborted)
-       we need to reinitialize object. */
-    if (mbedtls_test_driver_pake_hooks.forced_status == PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
-        TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
-                   PSA_SUCCESS);
-    }
-
-    /* --- psa_pake_set_role --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER),
-               expected_status_set_role);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* --- psa_pake_set_password_key --- */
-    /* psa_pake_set_password_key is dispatched by location and other
-       functions are dispatched by operation id (set during setup).
-       In case of dispatching by location fallback is performed when
-       transparent accelerators are not supported. */
-    if (forced_status_setup == PSA_ERROR_NOT_SUPPORTED) {
-        mbedtls_test_driver_pake_hooks.forced_status = PSA_ERROR_NOT_SUPPORTED;
-        expected_hits = 1;
-    } else {
-        mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    }
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_set_password_key(&operation, key),
-               expected_status_set_password);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* Restore expected_hits for next tests. */
-    if (forced_status_setup == PSA_ERROR_NOT_SUPPORTED) {
-        expected_hits = 0;
-    }
-
-    /* --- psa_pake_input --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
-                              input_buffer, size_key_share),
-               expected_status_input);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* --- psa_pake_abort --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_abort(&operation), expected_status_abort);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    /* --- psa_pake_output --- */
-    /* We need to setup pake again */
-    mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
-    TEST_EQUAL(psa_pake_abort(&operation), PSA_SUCCESS);
-
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
     TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
                PSA_SUCCESS);
 
@@ -3199,35 +3054,98 @@
     TEST_EQUAL(psa_pake_set_password_key(&operation, key),
                PSA_SUCCESS);
 
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
 
-    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;
+    /* Computation stage (driver entry points) */
+
+    switch (fut) {
+        case 0: /* setup (via input) */
+            /* --- 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;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       expected_status_setup);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
+            break;
+
+        case 1: /* setup (via output) */
+            /* --- 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;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       expected_status_setup);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
+            break;
+
+        case 2: /* input */
+            /* --- 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;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       expected_status_input);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, in_driver ? 2 : 1);
+            break;
+
+        case 3: /* output */
+            /* --- psa_pake_input (driver: setup, output) --- */
+            mbedtls_test_driver_pake_hooks.forced_setup_status = forced_status_setup;
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            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;
+            }
+            TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                       output_buffer, output_size, &output_len),
+                       expected_status_output);
+
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, in_driver ? 2 : 1);
+            if (forced_output->len > 0) {
+                TEST_EQUAL(output_len, forced_output->len);
+                TEST_EQUAL(memcmp(output_buffer, forced_output->x, output_len), 0);
+            }
+            break;
+
+        case 4: /* get_implicit_key */
+            /* Call driver setup indirectly */
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       PSA_SUCCESS);
+
+            /* --- psa_pake_get_implicit_key --- */
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            TEST_EQUAL(psa_pake_get_implicit_key(&operation, &implicit_key),
+                       expected_status_get_key);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
+
+            break;
+
+        case 5: /* abort */
+            /* Call driver setup indirectly */
+            TEST_EQUAL(psa_pake_input(&operation, PSA_PAKE_STEP_KEY_SHARE,
+                                      input_buffer, size_key_share),
+                       PSA_SUCCESS);
+
+            /* --- psa_pake_abort --- */
+            mbedtls_test_driver_pake_hooks.forced_status = forced_status;
+            mbedtls_test_driver_pake_hooks.hits = 0;
+            TEST_EQUAL(psa_pake_abort(&operation), expected_status_abort);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
+            break;
+
+        default:
+            break;
     }
 
-    TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
-                               output_buffer, output_size, &output_len),
-               expected_status_output);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
-    if (forced_output->len > 0) {
-        TEST_EQUAL(output_len, forced_output->len);
-        TEST_EQUAL(memcmp(output_buffer, forced_output->x, output_len), 0);
-    }
-
-    /* --- psa_pake_get_implicit_key --- */
-    mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-    mbedtls_test_driver_pake_hooks.hits = 0;
-
-    TEST_EQUAL(psa_pake_get_implicit_key(&operation, &implicit_key),
-               expected_status_get_key);
-
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, expected_hits);
-
     /* Clean up */
+    mbedtls_test_driver_pake_hooks.forced_setup_status = PSA_SUCCESS;
     mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
     TEST_EQUAL(psa_pake_abort(&operation), PSA_SUCCESS);
 exit:
@@ -3265,7 +3183,12 @@
     pake_in_driver = in_driver;
     mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
     mbedtls_test_driver_pake_hooks.hits = 0;
-    pake_expected_hit_count = 1;
+    /* driver setup is called indirectly through pake_output/pake_input */
+    if (pake_in_driver) {
+        pake_expected_hit_count = 2;
+    } else {
+        pake_expected_hit_count = 1;
+    }
 
     PSA_INIT();
 
@@ -3293,49 +3216,23 @@
     }
 
     if (!pake_in_driver) {
-        mbedtls_test_driver_pake_hooks.forced_status = PSA_ERROR_NOT_SUPPORTED;
+        mbedtls_test_driver_pake_hooks.forced_setup_status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     PSA_ASSERT(psa_pake_setup(&server, &cipher_suite));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, pake_expected_hit_count++);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
     PSA_ASSERT(psa_pake_setup(&client, &cipher_suite));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, pake_expected_hit_count++);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
 
-    /* Restore forced status and adjust pake_expected_hit_count */
-    mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
-    if (!pake_in_driver) {
-        pake_expected_hit_count--;
-    }
 
     PSA_ASSERT(psa_pake_set_role(&server, PSA_PAKE_ROLE_SERVER));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
-               pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
     PSA_ASSERT(psa_pake_set_role(&client, PSA_PAKE_ROLE_CLIENT));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
-               pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
-
-    /* psa_pake_set_password_key is dispatched by location and other
-       functions are dispatched by operation id (set during setup).
-       In case of dispatching by location fallback is performed when
-       transparent accelerators are not supported. We need to also adjust
-       expected hit counter. */
-    if (!pake_in_driver) {
-        mbedtls_test_driver_pake_hooks.forced_status = PSA_ERROR_NOT_SUPPORTED;
-        pake_expected_hit_count++;
-    }
-
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
-               pake_expected_hit_count++);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
     PSA_ASSERT(psa_pake_set_password_key(&client, key));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
-               pake_expected_hit_count++);
-
-    /* Restore forced status and adjust pake_expected_hit_count */
-    mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
-    if (!pake_in_driver) {
-        pake_expected_hit_count--;
-    }
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
 
     /* First round */
     ecjpake_do_round(alg, primitive_arg, &server, &client,