diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index f79d9b3..23d827e 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -19,6 +19,7 @@
 #define HAVE_RAM_AVAILABLE_128K
 
 #include "psa/crypto.h"
+#include "psa_crypto_slot_management.h"
 
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
@@ -245,7 +246,13 @@
 
     /* Persistence */
     if( lifetime == PSA_KEY_LIFETIME_VOLATILE )
-        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) == 0 );
+    {
+        TEST_ASSERT(
+            ( PSA_KEY_ID_VOLATILE_MIN <=
+              MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
+            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <=
+              PSA_KEY_ID_VOLATILE_MAX ) );
+    }
     else
     {
         TEST_ASSERT(
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 a1d542d..f22e6b7 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -972,6 +972,7 @@
     psa_set_key_bits( &attributes,
                       PSA_BYTES_TO_BITS( sizeof( key_material ) ) );
     psa_set_key_slot_number( &attributes, min_slot );
+    psa_set_key_id( &attributes, handle );
     if( ! check_key_attributes( handle, &attributes ) )
         goto exit;