diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index d1ec001..d6e9e14 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -268,37 +268,9 @@
                 return( PSA_SUCCESS );
             }
 #endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
-#ifdef MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
             *key_buffer_size = mbedtls_test_size_function( key_type, key_bits );
-            return( PSA_SUCCESS );
-#else /* MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */
-            if( PSA_KEY_TYPE_IS_KEY_PAIR( key_type ) )
-            {
-                int public_key_overhead =
-                    ( ( MBEDTLS_TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY == 1 )
-                      ? PSA_EXPORT_KEY_OUTPUT_SIZE( key_type, key_bits ) : 0 );
-                *key_buffer_size = MBEDTLS_TEST_DRIVER_KEY_CONTEXT_BASE_SIZE +
-                    MBEDTLS_TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE +
-                    public_key_overhead;
-            }
-            else if( PSA_KEY_TYPE_IS_PUBLIC_KEY( key_type ) )
-            {
-                *key_buffer_size = MBEDTLS_TEST_DRIVER_KEY_CONTEXT_BASE_SIZE +
-                    MBEDTLS_TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE;
-            }
-            else if ( !PSA_KEY_TYPE_IS_KEY_PAIR( key_type ) &&
-                      !PSA_KEY_TYPE_IS_PUBLIC_KEY ( key_type ) )
-            {
-                *key_buffer_size = MBEDTLS_TEST_DRIVER_KEY_CONTEXT_BASE_SIZE +
-                    ( MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR *
-                      ( ( key_bits + 7 ) / 8 ) );
-            }
-            else
-            {
-                return( PSA_ERROR_NOT_SUPPORTED );
-            }
-            return( PSA_SUCCESS );
-#endif /* MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */
+            return( ( *key_buffer_size != 0 ) ?
+                    PSA_SUCCESS : PSA_ERROR_NOT_SUPPORTED );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
         default:
diff --git a/tests/include/test/drivers/size.h b/tests/include/test/drivers/size.h
index 577e17b..b2665bd 100644
--- a/tests/include/test/drivers/size.h
+++ b/tests/include/test/drivers/size.h
@@ -29,68 +29,9 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
 #include <psa/crypto_driver_common.h>
 
-typedef struct {
-    unsigned int context;
-} mbedtls_test_driver_key_context_t;
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_BASE_SIZE
- *
- * 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 MBEDTLS_TEST_DRIVER_KEY_CONTEXT_BASE_SIZE \
-    sizeof( mbedtls_test_driver_key_context_t )
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_KEY_PAIR_SIZE
- *
- * Number of bytes included in every key context for a key pair.
- *
- * This pair size is for an ECC 256-bit private/public key pair.
- * Based on this value, the size of the private key can be derived by
- * subtracting the public key size below from this one.
- */
-
-#define MBEDTLS_TEST_DRIVER_KEY_CONTEXT_KEY_PAIR_SIZE      65
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE
- *
- * Number of bytes included in every key context for a public key.
- *
- * For ECC public keys, it needs 257 bits so 33 bytes.
- */
-#define MBEDTLS_TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE    33
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR
- *
- * Every key context for a symmetric key includes this many times the key size.
- */
-#define MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR   0
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY
- *
- * If this is true for a key pair, the key context includes space for the public key.
- * If this is false, no additional space is added for the public key.
- *
- * For this instance, store the public key with the private one.
- */
-#define MBEDTLS_TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY   1
-
-/** \def MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
- *
- * If MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION is defined, the test driver
- * provides a size_function entry point, otherwise, it does not.
- *
- * Some opaque drivers have the need to support a custom size for the storage
- * of key and context information. The size_function provides the ability to
- * provide that customization.
- */
-//#define MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
-
-#ifdef MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION
 size_t mbedtls_test_size_function(
     const psa_key_type_t key_type,
     const size_t key_bits );
-#endif /* MBEDTLS_TEST_DRIVER_KEY_CONTEXT_SIZE_FUNCTION */
 
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_TEST_DRIVERS_SIZE_H */
diff --git a/tests/src/drivers/test_driver_size.c b/tests/src/drivers/test_driver_size.c
index fd10209..d8bcaee 100644
--- a/tests/src/drivers/test_driver_size.c
+++ b/tests/src/drivers/test_driver_size.c
@@ -27,16 +27,75 @@
 #if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
 
 #include "test/drivers/size.h"
+#include "psa/crypto.h"
 
-#ifdef MBEDTLS_TEST_KEY_CONTEXT_SIZE_FUNCTION
+typedef struct {
+    unsigned int context;
+} test_driver_key_context_t;
+
+/*
+ * 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 )
+
+/*
+ * Number of bytes included in every key context for a key pair.
+ *
+ * This pair size is for an ECC 256-bit private/public key pair.
+ * Based on this value, the size of the private key can be derived by
+ * subtracting the public key size below from this one.
+ */
+#define TEST_DRIVER_KEY_CONTEXT_KEY_PAIR_SIZE      65
+
+/*
+ * Number of bytes included in every key context for a public key.
+ *
+ * For ECC public keys, it needs 257 bits so 33 bytes.
+ */
+#define TEST_DRIVER_KEY_CONTEXT_PUBLIC_KEY_SIZE    33
+
+/*
+ * Every key context for a symmetric key includes this many times the key size.
+ */
+#define TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR   0
+
+/*
+ * If this is true for a key pair, the key context includes space for the public key.
+ * If this is false, no additional space is added for the public key.
+ *
+ * For this instance, store the public key with the private one.
+ */
+#define TEST_DRIVER_KEY_CONTEXT_STORE_PUBLIC_KEY   1
+
 size_t mbedtls_test_size_function(
     const psa_key_type_t key_type,
     const size_t key_bits )
 {
-    (void) key_type;
-    (void) key_bits;
-    return 0;
-}
-#endif /*MBEDTLS_TEST_KEY_CONTEXT_SIZE_FUNCTION */
+    size_t key_buffer_size = 0;
 
+    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 );
+        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( key_type ) )
+    {
+        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 ( key_type ) )
+    {
+        key_buffer_size = TEST_DRIVER_KEY_CONTEXT_BASE_SIZE +
+                          ( TEST_DRIVER_KEY_CONTEXT_SYMMETRIC_FACTOR *
+                            ( ( key_bits + 7 ) / 8 ) );
+    }
+
+    return( key_buffer_size );
+}
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */