diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index f905ea4..f19f559 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -279,10 +279,19 @@
 #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_KEY_EXPORT_MAX_SIZE( key_type, key_bits ) : 0 );
+                *expected_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 ) )
+            {
                 *expected_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
                                  + TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE;
             }
-            else if( PSA_KEY_TYPE_IS_PUBLIC_KEY( attributes->core.type ) )
+            else if ( !PSA_KEY_TYPE_IS_KEY_PAIR( key_type ) &&
+                      !PSA_KEY_TYPE_IS_PUBLIC_KEY ( attributes->core.type ) )
             {
                 *expected_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
                                  + TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR
@@ -300,7 +309,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
     }
 }
-#endif /* PSA_CRYPTO_DRIVER_PRESENT */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
 
 psa_status_t psa_driver_wrapper_generate_key( const psa_key_attributes_t *attributes,
                                               psa_key_slot_t *slot )
diff --git a/tests/include/test/drivers/size.h b/tests/include/test/drivers/size.h
index 831adbb..4bfe986 100644
--- a/tests/include/test/drivers/size.h
+++ b/tests/include/test/drivers/size.h
@@ -35,8 +35,8 @@
 
 /** \def TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
  *
- * This macro returns the base size for the key context. It should include
- * the size for any driver context information stored with each key.
+ * This macro returns the base size for the key context. It is the size of the
+ * driver specific information stored in each key context.
  */
 #define TEST_DRIVER_KEY_CONTEXT_BASE_SIZE          sizeof( test_driver_key_context_t )
 
@@ -92,4 +92,4 @@
 #endif /* TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */
 
 #endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_TEST_DRIVERS_KEYGEN_H */
+#endif /* PSA_CRYPTO_TEST_DRIVERS_SIZE_H */
diff --git a/tests/src/drivers/size.c b/tests/src/drivers/size.c
index 05f8a98..16a8692 100644
--- a/tests/src/drivers/size.c
+++ b/tests/src/drivers/size.c
@@ -25,14 +25,9 @@
 #endif
 
 #if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
-#include "psa/crypto.h"
-#include "psa_crypto_core.h"
-#include "mbedtls/error.h"
 
 #include "test/drivers/size.h"
 
-#include <string.h>
-
 #ifdef TEST_KEY_CONTEXT_SIZE_FUNCTION
 size_t test_size_function(
     const psa_key_type_t key_type,