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 11b8866..79d658f 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -1467,6 +1467,7 @@
             ( validate > 0 ? PSA_SUCCESS : PSA_ERROR_NOT_PERMITTED );
     }
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_register_se_driver( MIN_DRIVER_LOCATION, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -1496,6 +1497,7 @@
     PSA_ASSERT( psa_purge_key( id ) );
 
     /* Restart and try again. */
+    mbedtls_test_set_step( 2 );
     PSA_SESSION_DONE( );
     PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index 7d3c7a8..bafb7d8 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -123,6 +123,7 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Import a key. */
@@ -169,6 +170,7 @@
     psa_reset_key_attributes( &attributes );
 
     /* Do something that invalidates the key. */
+    mbedtls_test_set_step( 2 );
     if( ! invalidate_key( invalidate_method, key ) )
         goto exit;
     if( ! invalidate_psa( invalidate_method ) )
@@ -218,6 +220,7 @@
 
     TEST_USES_KEY_ID( id );
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -267,6 +270,7 @@
      * Do something that wipes key data in volatile memory or destroy the
      * key.
      */
+    mbedtls_test_set_step( 2 );
     if( ! invalidate_key( invalidate_method, id ) )
         goto exit;
     if( ! invalidate_psa( invalidate_method ) )