diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 9440195cc8..07888f1de2 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -239,61 +239,56 @@ psa_status_t psa_driver_wrapper_verify_hash( psa_key_slot_t *slot,
 }
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-/** Calculate the size to allocate for buffering a key with given attributes.
+/** Get the key buffer size for the key material of a generated key in the
+ *  case of an opaque driver without storage.
  *
- * This function provides a way to get the expected size for storing a key with
- * the given attributes. This will be the size of the export representation for
- * cleartext keys, and a driver-defined size for keys stored by opaque drivers.
- *
- * \param[in] attributes        The key attribute structure of the key to store.
- * \param[out] expected_size    On success, a byte size large enough to contain
- *                              the declared key.
+ * \param[in] attributes  The key attributes.
+ * \param[out] key_buffer_size  Minimum buffer size to contain the key material
  *
  * \retval #PSA_SUCCESS
+ *         The minimum size for a buffer to contain the key material has been
+ *         returned successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The size in bits of the key is not valid.
  * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         The type and/or the size in bits of the key or the combination of
+ *         the two is not supported.
  */
-static psa_status_t get_expected_key_size( const psa_key_attributes_t *attributes,
-                                           size_t *expected_size )
+static psa_status_t get_key_buffer_size(
+    const psa_key_attributes_t *attributes,
+    size_t *key_buffer_size )
 {
-    size_t buffer_size = 0;
     psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime );
     psa_key_type_t key_type = attributes->core.type;
     size_t key_bits = attributes->core.bits;
 
+    *key_buffer_size = 0;
     switch( location )
     {
-        case PSA_KEY_LOCATION_LOCAL_STORAGE:
-            buffer_size = PSA_EXPORT_KEY_OUTPUT_SIZE( key_type, key_bits );
-
-            if( buffer_size == 0 )
-                return( PSA_ERROR_NOT_SUPPORTED );
-
-            *expected_size = buffer_size;
-            return( PSA_SUCCESS );
-
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
 #ifdef TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
-            *expected_size = test_size_function( key_type, key_bits );
+            *key_buffer_size = test_size_function( key_type, key_bits );
             return( PSA_SUCCESS );
 #else /* TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */
             if( PSA_KEY_TYPE_IS_KEY_PAIR( key_type ) )
             {
-                int public_key_overhead = ( ( TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY == 1 ) ?
-                                           PSA_EXPORT_KEY_OUTPUT_SIZE( key_type, key_bits ) : 0 );
-                *expected_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
+                int public_key_overhead =
+                    ( ( TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY == 1 ) ?
+                      PSA_EXPORT_KEY_OUTPUT_SIZE( key_type, key_bits ) : 0 );
+                *key_buffer_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
                                  + TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE
                                  + public_key_overhead;
             }
-            else if( PSA_KEY_TYPE_IS_PUBLIC_KEY( attributes->core.type ) )
+            else if( PSA_KEY_TYPE_IS_PUBLIC_KEY( key_type ) )
             {
-                *expected_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
+                *key_buffer_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
                                  + TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE;
             }
             else if ( !PSA_KEY_TYPE_IS_KEY_PAIR( key_type ) &&
-                      !PSA_KEY_TYPE_IS_PUBLIC_KEY ( attributes->core.type ) )
+                      !PSA_KEY_TYPE_IS_PUBLIC_KEY ( key_type ) )
             {
-                *expected_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
+                *key_buffer_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
                                  + TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR
                                  * ( ( key_bits + 7 ) / 8 );
             }
@@ -345,10 +340,21 @@ psa_status_t psa_driver_wrapper_generate_key( const psa_key_attributes_t *attrib
      * with storage. */
     if( slot->key.data == NULL )
     {
-        status = get_expected_key_size( attributes, &export_size );
-        if( status != PSA_SUCCESS )
-            return( status );
+        if( location == PSA_KEY_LOCATION_LOCAL_STORAGE )
+        {
+            export_size = PSA_KEY_EXPORT_MAX_SIZE( attributes->core.type,
+                                                   attributes->core.bits );
 
+            if( export_size == 0 )
+                return( PSA_ERROR_NOT_SUPPORTED );
+        }
+        else
+        {
+            status = get_key_buffer_size( attributes, &export_size );
+            if( status != PSA_SUCCESS )
+                return( status );
+        }
+    
         slot->key.data = mbedtls_calloc(1, export_size);
         if( slot->key.data == NULL )
             return( PSA_ERROR_INSUFFICIENT_MEMORY );