diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 3e356f9..9881eae 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -45,6 +45,9 @@
 
 /** Purge the in-memory cache of persistent keys recorded with
  * #TEST_USES_KEY_ID.
+ *
+ * Call this function before calling PSA_DONE() if it's ok for
+ * persistent keys to still exist at this point.
  */
 void mbedtls_test_psa_purge_key_cache( void );
 
@@ -79,6 +82,8 @@
 #else /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
 
 #define TEST_USES_KEY_ID( key_id ) ( (void) ( key_id ) )
+#define mbedtls_test_psa_purge_key_storage( ) ( (void) 0 )
+#define mbedtls_test_psa_purge_key_cache( ) ( (void) 0 )
 
 #endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
 
@@ -108,13 +113,36 @@
     }                                                                   \
     while( 0 )
 
-/** Shut down the PSA Crypto subsystem. Expect a clean shutdown, with no slots
- * in use.
+/** Shut down the PSA Crypto subsystem and destroy persistent keys.
+ * Expect a clean shutdown, with no slots in use.
+ *
+ * If some key slots are still in use, record the test case as failed,
+ * but continue executing. This macro is suitable (and primarily intended)
+ * for use in the cleanup section of test functions.
+ *
+ * \note Persistent keys must be recorded with #TEST_USES_KEY_ID before
+ *       creating them.
  */
 #define PSA_DONE( )                                                     \
     do                                                                  \
     {                                                                   \
         test_fail_if_psa_leaking( __LINE__, __FILE__ );                 \
+        mbedtls_test_psa_purge_key_storage( );                          \
+        mbedtls_psa_crypto_free( );                                     \
+    }                                                                   \
+    while( 0 )
+
+/** Shut down the PSA Crypto subsystem, allowing persistent keys to survive.
+ * Expect a clean shutdown, with no slots in use.
+ *
+ * If some key slots are still in use, record the test case as failed and
+ * jump to the `exit` label.
+ */
+#define PSA_SESSION_DONE( )                                             \
+    do                                                                  \
+    {                                                                   \
+        mbedtls_test_psa_purge_key_cache( );                            \
+        ASSERT_PSA_PRISTINE( );                                         \
         mbedtls_psa_crypto_free( );                                     \
     }                                                                   \
     while( 0 )
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index be7c246..11b8866 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -769,10 +769,10 @@
 
 static void psa_purge_storage( void )
 {
+    /* The generic code in mbedtls_test_psa_purge_key_storage()
+     * (which is called by PSA_DONE()) doesn't take care of things that are
+     * specific to dynamic secure elements. */
     psa_key_location_t location;
-
-    mbedtls_test_psa_purge_key_storage( );
-
     /* Purge the transaction file. */
     psa_crypto_stop_transaction( );
     /* Purge driver persistent data. */
@@ -1496,7 +1496,7 @@
     PSA_ASSERT( psa_purge_key( id ) );
 
     /* Restart and try again. */
-    PSA_DONE( );
+    PSA_SESSION_DONE( );
     PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
     if( ! check_key_attributes( id, &attributes ) )
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index 7c55c71..7d3c7a8 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -86,8 +86,7 @@
         case INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN:
         case INVALIDATE_BY_PURGING_WITH_SHUTDOWN:
             /* All keys must have been closed. */
-            mbedtls_test_psa_purge_key_cache( );
-            PSA_DONE( );
+            PSA_SESSION_DONE( );
             break;
         case INVALIDATE_BY_SHUTDOWN:
             /* Some keys may remain behind, and we're testing that this
@@ -339,7 +338,6 @@
     psa_reset_key_attributes( &read_attributes );
 
     PSA_DONE( );
-    mbedtls_test_psa_purge_key_storage( );
     mbedtls_free( reexported );
 }
 /* END_CASE */
@@ -413,7 +411,6 @@
     psa_reset_key_attributes( &attributes );
 
     PSA_DONE( );
-    mbedtls_test_psa_purge_key_storage( );
 }
 /* END_CASE */
 
@@ -472,9 +469,6 @@
 
 exit:
     PSA_DONE( );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    mbedtls_test_psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */
 
@@ -610,9 +604,6 @@
 
     PSA_DONE( );
     mbedtls_free( export_buffer );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    mbedtls_test_psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */
 
@@ -728,9 +719,6 @@
 
     PSA_DONE( );
     mbedtls_free( export_buffer );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    mbedtls_test_psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */