diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 17dfe5f..2a8fede 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -442,4 +442,24 @@
 #define AES_PSA_DONE() ((void) 0)
 #endif /* MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO */
 
+#if !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) &&                        \
+    defined(MBEDTLS_CTR_DRBG_C) &&                                      \
+    defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
+/* When AES_C is not defined and PSA does not have an external RNG,
+ * then CTR_DRBG uses PSA to perform AES-ECB. In this scenario 1 key
+ * slot is used internally from PSA to hold the AES key and it should
+ * not be taken into account when evaluating remaining open slots. */
+#define MBEDTLS_TEST_PSA_INTERNAL_KEYS_FOR_DRBG 1
+#else
+#define MBEDTLS_TEST_PSA_INTERNAL_KEYS_FOR_DRBG 0
+#endif
+
+/** The number of volatile keys that PSA crypto uses internally.
+ *
+ * We expect that many volatile keys to be in use after a successful
+ * psa_crypto_init().
+ */
+#define MBEDTLS_TEST_PSA_INTERNAL_KEYS          \
+    MBEDTLS_TEST_PSA_INTERNAL_KEYS_FOR_DRBG
+
 #endif /* PSA_CRYPTO_HELPERS_H */
diff --git a/tests/src/psa_crypto_helpers.c b/tests/src/psa_crypto_helpers.c
index 1581eec..1069edd 100644
--- a/tests/src/psa_crypto_helpers.c
+++ b/tests/src/psa_crypto_helpers.c
@@ -74,21 +74,9 @@
 
     mbedtls_psa_get_stats(&stats);
 
-#if !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) &&                        \
-    defined(MBEDTLS_CTR_DRBG_C) &&                                      \
-    defined(MBEDTLS_CTR_DRBG_USE_PSA_CRYPTO)
-    /* When AES_C is not defined and PSA does not have an external RNG,
-     * then CTR_DRBG uses PSA to perform AES-ECB. In this scenario 1 key
-     * slot is used internally from PSA to hold the AES key and it should
-     * not be taken into account when evaluating remaining open slots. */
     if (stats.volatile_slots > 1) {
         return "A volatile slot has not been closed properly.";
     }
-#else
-    if (stats.volatile_slots != 0) {
-        return "A volatile slot has not been closed properly.";
-    }
-#endif
     if (stats.persistent_slots != 0) {
         return "A persistent slot has not been closed properly.";
     }
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_init.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_init.function
index 9ff33a6..2fd282e 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_init.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_init.function
@@ -8,6 +8,23 @@
 #include "mbedtls/entropy.h"
 #include "entropy_poll.h"
 
+static int check_stats(void)
+{
+    mbedtls_psa_stats_t stats;
+    mbedtls_psa_get_stats(&stats);
+
+    TEST_EQUAL(stats.volatile_slots, MBEDTLS_TEST_PSA_INTERNAL_KEYS);
+    TEST_EQUAL(stats.persistent_slots, 0);
+    TEST_EQUAL(stats.external_slots, 0);
+    TEST_EQUAL(stats.half_filled_slots, 0);
+    TEST_EQUAL(stats.locked_slots, 0);
+
+    return 1;
+
+exit:
+    return 0;
+}
+
 #define ENTROPY_MIN_NV_SEED_SIZE                                        \
     MAX(MBEDTLS_ENTROPY_MIN_PLATFORM, MBEDTLS_ENTROPY_BLOCK_SIZE)
 
@@ -187,10 +204,19 @@
     psa_status_t status;
     int i;
     for (i = 0; i < count; i++) {
+        mbedtls_test_set_step(2 * i);
         status = psa_crypto_init();
         PSA_ASSERT(status);
+        if (!check_stats()) {
+            goto exit;
+        }
+
+        mbedtls_test_set_step(2 * i);
         status = psa_crypto_init();
         PSA_ASSERT(status);
+        if (!check_stats()) {
+            goto exit;
+        }
         PSA_DONE();
     }
 }
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
index 742f9b1..af3b946 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
@@ -225,11 +225,8 @@
 invalid handle: huge
 invalid_handle:INVALID_HANDLE_HUGE:PSA_ERROR_INVALID_HANDLE
 
-Key slot count: less than maximum
-many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT - 1
-
 Key slot count: maximum
-many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT
+many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT - MBEDTLS_TEST_PSA_INTERNAL_KEYS
 
 Key slot count: try to overfill, destroy first
 fill_key_store:0