diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index 003a35d64..b2529ce08 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -30,6 +30,10 @@ typedef struct {
     /* Count the amount of times one of the key management driver functions
      * is called. */
     unsigned long hits;
+    /* Subset of hits which only counts public key export operations */
+    unsigned long hits_export_public_key;
+    /* Subset of hits which only counts key generation operations */
+    unsigned long hits_generate_key;
     /* Location of the last key management driver called to import a key. */
     psa_key_location_t location;
 } mbedtls_test_driver_key_management_hooks_t;
@@ -38,7 +42,7 @@ typedef struct {
  * sense that no PSA specification will assign a meaning to this location
  * (stated first in version 1.0.1 of the specification) and that it is not
  * used as a location of an opaque test drivers. */
-#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0x800000 }
+#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0, 0, 0x800000 }
 static inline mbedtls_test_driver_key_management_hooks_t
 mbedtls_test_driver_key_management_hooks_init(void)
 {
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index 4d06d23b8..e6d7a0107 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -93,6 +93,7 @@ psa_status_t mbedtls_test_transparent_generate_key(
     uint8_t *key, size_t key_size, size_t *key_length)
 {
     ++mbedtls_test_driver_key_management_hooks.hits;
+    ++mbedtls_test_driver_key_management_hooks.hits_generate_key;
 
     if (mbedtls_test_driver_key_management_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_key_management_hooks.forced_status;
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index b589d81bc..a58f0835b 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -238,6 +238,7 @@ generate_key through transparent driver: fake
 generate_key:PSA_SUCCESS:"ab45435712649cb30bbddac49197eebf2740ffc7f874d9244c3460f54f322d3a":PSA_SUCCESS
 
 generate_key through transparent driver: in-driver
+depends_on:MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR
 generate_key:PSA_SUCCESS:"":PSA_SUCCESS
 
 generate_key through transparent driver: fallback
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 0394735aa..64b3d1b67 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -300,13 +300,14 @@ void generate_key(int force_status_arg,
                 fake_output->len;
     }
 
-    mbedtls_test_driver_key_management_hooks.hits = 0;
-    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
-
     PSA_ASSERT(psa_crypto_init());
 
+    mbedtls_test_driver_key_management_hooks.hits = 0;
+    mbedtls_test_driver_key_management_hooks.hits_generate_key = 0;
+    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
+
     actual_status = psa_generate_key(&attributes, &key);
-    TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits_generate_key, 1);
     TEST_EQUAL(actual_status, expected_status);
 
     if (actual_status == PSA_SUCCESS) {