diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 6932318..1c348e8 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1500,7 +1500,7 @@
 {
     psa_status_t status;
 
-    if( attributes->core.lifetime != PSA_KEY_LIFETIME_VOLATILE )
+    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( attributes->core.lifetime ) )
     {
         status = psa_validate_persistent_key_parameters(
             attributes->core.lifetime, attributes->core.id,
@@ -1660,7 +1660,7 @@
     (void) driver;
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( slot->attr.lifetime != PSA_KEY_LIFETIME_VOLATILE )
+    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( slot->attr.lifetime ) )
     {
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
         if( driver != NULL )
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 6cd6a11..193959a 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -200,7 +200,10 @@
     }
     else
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    if( lifetime != PSA_KEY_LIFETIME_PERSISTENT )
+    if( ( PSA_KEY_LIFETIME_GET_LOCATION( lifetime )
+            != PSA_KEY_LOCATION_LOCAL_STORAGE ) ||
+        ( PSA_KEY_LIFETIME_GET_PERSISTENCE( lifetime )
+            != PSA_KEY_PERSISTENCE_DEFAULT ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 472253d..db2aa96 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -88,8 +88,8 @@
  */
 static inline int psa_key_lifetime_is_external( psa_key_lifetime_t lifetime )
 {
-    return( lifetime != PSA_KEY_LIFETIME_VOLATILE &&
-            lifetime != PSA_KEY_LIFETIME_PERSISTENT );
+    return( PSA_KEY_LIFETIME_GET_LOCATION( lifetime )
+                != PSA_KEY_LOCATION_LOCAL_STORAGE );
 }
 
 /** Test whether the given parameters are acceptable for a persistent key.