diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index 8030724..99ca8f2 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -33,7 +33,14 @@
        pake_output/pake_input (added to distinguish forced statuses). */
     psa_status_t forced_setup_status;
     /* Count the amount of times PAKE driver functions are called. */
-    unsigned long hits;
+    struct {
+        unsigned long total;
+        unsigned long setup;
+        unsigned long input;
+        unsigned long output;
+        unsigned long implicit_key;
+        unsigned long abort;
+    } hits;
     /* Status returned by the last PAKE driver function call. */
     psa_status_t driver_status;
     /* Output returned by pake_output */
@@ -41,7 +48,7 @@
     size_t forced_output_length;
 } mbedtls_test_driver_pake_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_PAKE_INIT { PSA_SUCCESS, PSA_SUCCESS, 0, PSA_SUCCESS, NULL, 0 }
+#define MBEDTLS_TEST_DRIVER_PAKE_INIT { PSA_SUCCESS, PSA_SUCCESS, {0, 0, 0, 0, 0, 0}, PSA_SUCCESS, NULL, 0 }
 static inline mbedtls_test_driver_pake_hooks_t
 mbedtls_test_driver_pake_hooks_init(void)
 {
diff --git a/tests/src/drivers/test_driver_pake.c b/tests/src/drivers/test_driver_pake.c
index 03f387f..7eafe14 100644
--- a/tests/src/drivers/test_driver_pake.c
+++ b/tests/src/drivers/test_driver_pake.c
@@ -37,7 +37,8 @@
     mbedtls_transparent_test_driver_pake_operation_t *operation,
     const psa_crypto_driver_pake_inputs_t *inputs)
 {
-    mbedtls_test_driver_pake_hooks.hits++;
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.setup++;
 
     if (mbedtls_test_driver_pake_hooks.forced_setup_status != PSA_SUCCESS) {
         mbedtls_test_driver_pake_hooks.driver_status =
@@ -69,7 +70,8 @@
     size_t output_size,
     size_t *output_length)
 {
-    mbedtls_test_driver_pake_hooks.hits++;
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.output++;
 
     if (mbedtls_test_driver_pake_hooks.forced_output != NULL) {
         if (output_size < mbedtls_test_driver_pake_hooks.forced_output_length) {
@@ -116,7 +118,8 @@
     const uint8_t *input,
     size_t input_length)
 {
-    mbedtls_test_driver_pake_hooks.hits++;
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.input++;
 
     if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
         mbedtls_test_driver_pake_hooks.driver_status =
@@ -147,7 +150,8 @@
     mbedtls_transparent_test_driver_pake_operation_t *operation,
     uint8_t *output, size_t output_size, size_t *output_length)
 {
-    mbedtls_test_driver_pake_hooks.hits++;
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.implicit_key++;
 
     if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
         mbedtls_test_driver_pake_hooks.driver_status =
@@ -177,7 +181,8 @@
 psa_status_t mbedtls_test_transparent_pake_abort(
     mbedtls_transparent_test_driver_pake_operation_t *operation)
 {
-    mbedtls_test_driver_pake_hooks.hits++;
+    mbedtls_test_driver_pake_hooks.hits.total++;
+    mbedtls_test_driver_pake_hooks.hits.abort++;
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
         defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index fa7aa7b..d63371b 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -822,48 +822,41 @@
 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 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_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:0
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:0
 
 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_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:1
+pake_operations:"abcd":PSA_ERROR_GENERIC_ERROR:PSA_SUCCESS:"":PSA_ERROR_GENERIC_ERROR:1
 
 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
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR:2
 
 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
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR: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
+pake_operations:"abcd":PSA_SUCCESS:PSA_SUCCESS:"1234":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
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":PSA_ERROR_GENERIC_ERROR: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
+pake_operations:"abcd":PSA_SUCCESS:PSA_ERROR_GENERIC_ERROR:"":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_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:0
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:0
 
 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
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:1
 
 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
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:2
 
 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
+depends_on:!MBEDTLS_PSA_BUILTIN_PAKE
+pake_operations:"abcd":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS:"":PSA_ERROR_NOT_SUPPORTED:3
 
 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
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 0bc0a32..b0aac53 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -52,7 +52,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_g1_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_g1_len, expected_size_key_share);
             s_g1_off = buffer0_off;
@@ -60,7 +60,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x1_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_x1_pk_len, expected_size_zk_public);
             s_x1_pk_off = buffer0_off;
@@ -68,7 +68,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x1_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(s_x1_pr_len, max_expected_size_zk_proof);
             s_x1_pr_off = buffer0_off;
@@ -76,7 +76,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_g2_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_g2_len, expected_size_key_share);
             s_g2_off = buffer0_off;
@@ -84,7 +84,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x2_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_x2_pk_len, expected_size_zk_public);
             s_x2_pk_off = buffer0_off;
@@ -92,7 +92,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x2_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(s_x2_pr_len, max_expected_size_zk_proof);
             s_x2_pr_off = buffer0_off;
@@ -102,42 +102,42 @@
                 /* Client first round Input */
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_g1_off, s_g1_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x1_pk_off,
                                         s_x1_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x1_pr_off,
                                         s_x1_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_g2_off,
                                         s_g2_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x2_pk_off,
                                         s_x2_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x2_pr_off,
                                         s_x2_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
             }
@@ -149,7 +149,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_g1_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_g1_len, expected_size_key_share);
             c_g1_off = buffer1_off;
@@ -157,7 +157,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x1_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_x1_pk_len, expected_size_zk_public);
             c_x1_pk_off = buffer1_off;
@@ -165,7 +165,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x1_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(c_x1_pr_len, max_expected_size_zk_proof);
             c_x1_pr_off = buffer1_off;
@@ -173,7 +173,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_g2_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_g2_len, expected_size_key_share);
             c_g2_off = buffer1_off;
@@ -181,7 +181,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x2_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_x2_pk_len, expected_size_zk_public);
             c_x2_pk_off = buffer1_off;
@@ -189,7 +189,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x2_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(c_x2_pr_len, max_expected_size_zk_proof);
             c_x2_pr_off = buffer1_off;
@@ -199,42 +199,42 @@
                 /* Client first round Input */
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_g1_off, s_g1_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x1_pk_off,
                                         s_x1_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x1_pr_off,
                                         s_x1_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_g2_off,
                                         s_g2_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x2_pk_off,
                                         s_x2_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x2_pr_off,
                                         s_x2_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
             }
@@ -242,37 +242,37 @@
             /* Server first round Input */
             status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
                                     buffer1 + c_g1_off, c_g1_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                     buffer1 + c_x1_pk_off, c_x1_pk_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
                                     buffer1 + c_x1_pr_off, c_x1_pr_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
                                     buffer1 + c_g2_off, c_g2_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                     buffer1 + c_x2_pk_off, c_x2_pk_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
                                     buffer1 + c_x2_pr_off, c_x2_pr_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
@@ -285,7 +285,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_a_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_a_len, expected_size_key_share);
             s_a_off = buffer0_off;
@@ -293,7 +293,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x2s_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(s_x2s_pk_len, expected_size_zk_public);
             s_x2s_pk_off = buffer0_off;
@@ -301,7 +301,7 @@
             PSA_ASSERT(psa_pake_output(server, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer0 + buffer0_off,
                                        512 - buffer0_off, &s_x2s_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(s_x2s_pr_len, max_expected_size_zk_proof);
             s_x2s_pr_off = buffer0_off;
@@ -311,21 +311,21 @@
                 /* Client second round Input */
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_a_off, s_a_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x2s_pk_off,
                                         s_x2s_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x2s_pr_off,
                                         s_x2s_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
             }
@@ -336,7 +336,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_KEY_SHARE,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_a_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_a_len, expected_size_key_share);
             c_a_off = buffer1_off;
@@ -344,7 +344,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x2s_pk_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(c_x2s_pk_len, expected_size_zk_public);
             c_x2s_pk_off = buffer1_off;
@@ -352,7 +352,7 @@
             PSA_ASSERT(psa_pake_output(client, PSA_PAKE_STEP_ZK_PROOF,
                                        buffer1 + buffer1_off,
                                        512 - buffer1_off, &c_x2s_pr_len));
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_LE_U(c_x2s_pr_len, max_expected_size_zk_proof);
             c_x2s_pr_off = buffer1_off;
@@ -362,21 +362,21 @@
                 /* Client second round Input */
                 status = psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
                                         buffer0 + s_a_off, s_a_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PUBLIC,
                                         buffer0 + s_x2s_pk_off,
                                         s_x2s_pk_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
 
                 status = psa_pake_input(client, PSA_PAKE_STEP_ZK_PROOF,
                                         buffer0 + s_x2s_pr_off,
                                         s_x2s_pr_len);
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                            pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
                 TEST_EQUAL(status, PSA_SUCCESS);
             }
@@ -384,19 +384,19 @@
             /* Server second round Input */
             status = psa_pake_input(server, PSA_PAKE_STEP_KEY_SHARE,
                                     buffer1 + c_a_off, c_a_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PUBLIC,
                                     buffer1 + c_x2s_pk_off, c_x2s_pk_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
             status = psa_pake_input(server, PSA_PAKE_STEP_ZK_PROOF,
                                     buffer1 + c_x2s_pr_off, c_x2s_pr_len);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                        pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
             TEST_EQUAL(status, PSA_SUCCESS);
 
@@ -2976,70 +2976,57 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_JPAKE */
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
 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,
+                     data_t *forced_output, int expected_status_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_setup = expected_status_setup_arg;
-    psa_status_t expected_status_input = expected_status_input_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_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;
-    int in_driver = 1;
     psa_key_derivation_operation_t implicit_key =
         PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_pake_primitive_t primitive = 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();
     unsigned char *input_buffer = NULL;
-    const size_t size_key_share = PSA_PAKE_INPUT_SIZE(alg, primitive,
+    const size_t size_key_share = PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                                       PSA_PAKE_STEP_KEY_SHARE);
     unsigned char *output_buffer = NULL;
     size_t output_len = 0;
-    size_t output_size = PSA_PAKE_OUTPUT_SIZE(alg, primitive,
+    size_t output_size = PSA_PAKE_OUTPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                               PSA_PAKE_STEP_KEY_SHARE);
+    int in_driver = (forced_status_setup_arg == PSA_SUCCESS);
 
     ASSERT_ALLOC(input_buffer,
-                 PSA_PAKE_INPUT_SIZE(alg, primitive,
+                 PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                      PSA_PAKE_STEP_KEY_SHARE));
     memset(input_buffer, 0xAA, size_key_share);
 
     ASSERT_ALLOC(output_buffer,
-                 PSA_PAKE_INPUT_SIZE(alg, primitive,
+                 PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                      PSA_PAKE_STEP_KEY_SHARE));
     memset(output_buffer, 0x55, output_size);
 
-    if (forced_status_setup_arg == PSA_ERROR_NOT_SUPPORTED) {
-        in_driver = 0;
-    }
-
     PSA_INIT();
 
+    mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_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_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_import_key(&attributes, pw_data->x, pw_data->len,
                                   &key));
     }
 
-    psa_pake_cs_set_algorithm(&cipher_suite, alg);
+    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, hash_alg);
+    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
 
     mbedtls_test_driver_pake_hooks.forced_status = forced_status_setup;
 
@@ -3054,7 +3041,7 @@
     TEST_EQUAL(psa_pake_set_password_key(&operation, key),
                PSA_SUCCESS);
 
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
 
     /* Computation stage (driver entry points) */
 
@@ -3063,54 +3050,64 @@
             /* --- 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;
+            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_setup);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 1);
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
             break;
 
         case 1: /* setup (via output) */
-            /* --- psa_pake_input (driver: setup, input) --- */
+            /* --- 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;
-            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);
+            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);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
             break;
 
         case 2: /* input */
             /* --- 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;
-            mbedtls_test_driver_pake_hooks.hits = 0;
+            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_input);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, in_driver ? 3 : 1);
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 3 : 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.input, in_driver ? 1 : 0);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, in_driver ? 1 : 0);
             break;
 
         case 3: /* output */
-            /* --- psa_pake_input (driver: setup, output, (abort)) --- */
+            /* --- 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;
-            mbedtls_test_driver_pake_hooks.hits = 0;
+            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;
             }
             TEST_EQUAL(psa_pake_output(&operation, PSA_PAKE_STEP_KEY_SHARE,
                                        output_buffer, output_size, &output_len),
-                       expected_status_output);
+                       expected_status);
 
             if (forced_output->len > 0) {
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, in_driver ? 2 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 2 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.output, in_driver ? 1 : 0);
                 TEST_EQUAL(output_len, forced_output->len);
                 TEST_EQUAL(memcmp(output_buffer, forced_output->x, output_len), 0);
             } else {
-                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, in_driver ? 3 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, in_driver ? 3 : 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.setup, 1);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.output, in_driver ? 1 : 0);
+                TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, in_driver ? 1 : 0);
             }
             break;
 
@@ -3126,10 +3123,12 @@
 
             /* --- psa_pake_get_implicit_key --- */
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
-            mbedtls_test_driver_pake_hooks.hits = 0;
+            memset(&mbedtls_test_driver_pake_hooks.hits, 0, sizeof(mbedtls_test_driver_pake_hooks.hits));
             TEST_EQUAL(psa_pake_get_implicit_key(&operation, &implicit_key),
-                       expected_status_get_key);
-            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 2);
+                       expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 2);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.implicit_key, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, 1);
 
             break;
 
@@ -3141,9 +3140,10 @@
 
             /* --- 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);
+            memset(&mbedtls_test_driver_pake_hooks.hits, 0, sizeof(mbedtls_test_driver_pake_hooks.hits));
+            TEST_EQUAL(psa_pake_abort(&operation), expected_status);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 1);
+            TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.abort, 1);
             break;
 
         default:
@@ -3187,8 +3187,6 @@
     psa_key_derivation_operation_t client_derive =
         PSA_KEY_DERIVATION_OPERATION_INIT;
     pake_in_driver = in_driver;
-    mbedtls_test_driver_pake_hooks.forced_status = PSA_SUCCESS;
-    mbedtls_test_driver_pake_hooks.hits = 0;
     /* driver setup is called indirectly through pake_output/pake_input */
     if (pake_in_driver) {
         pake_expected_hit_count = 2;
@@ -3198,6 +3196,8 @@
 
     PSA_INIT();
 
+    mbedtls_test_driver_pake_hooks = mbedtls_test_driver_pake_hooks_init();
+
     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);
@@ -3226,19 +3226,19 @@
     }
 
     PSA_ASSERT(psa_pake_setup(&server, &cipher_suite));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_setup(&client, &cipher_suite));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
 
 
     PSA_ASSERT(psa_pake_set_role(&server, PSA_PAKE_ROLE_SERVER));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_set_role(&client, PSA_PAKE_ROLE_CLIENT));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_set_password_key(&client, key));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits, 0);
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
 
     /* First round */
     ecjpake_do_round(alg, primitive_arg, &server, &client,
@@ -3255,7 +3255,7 @@
     }
 
     PSA_ASSERT(psa_pake_get_implicit_key(&server, &server_derive));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
 
     /* After the key is obtained operation is aborted.
@@ -3265,7 +3265,7 @@
     }
 
     PSA_ASSERT(psa_pake_get_implicit_key(&client, &client_derive));
-    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
+    TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total,
                pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
 exit:
     psa_key_derivation_abort(&server_derive);