diff --git a/tests/src/drivers/keygen.c b/tests/src/drivers/keygen.c
index 4f30f0e..a21ec27 100644
--- a/tests/src/drivers/keygen.c
+++ b/tests/src/drivers/keygen.c
@@ -1,5 +1,6 @@
 /*
- * Test driver for signature functions
+ * Test driver for generating keys.
+ * Currently only supports generating ECC keys.
  */
 /*  Copyright The Mbed TLS Contributors
  *  SPDX-License-Identifier: Apache-2.0
@@ -63,11 +64,12 @@
 
     /* Copied from psa_crypto.c */
 #if defined(MBEDTLS_ECP_C)
-    if ( PSA_KEY_TYPE_IS_ECC( attributes->core.type ) && PSA_KEY_TYPE_IS_KEY_PAIR( attributes->core.type ) )
+    if ( PSA_KEY_TYPE_IS_ECC( psa_get_key_type( attributes ) )
+         && PSA_KEY_TYPE_IS_KEY_PAIR( psa_get_key_type( attributes ) ) )
     {
-        psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY( attributes->core.type );
+        psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY( psa_get_key_type( attributes ) );
         mbedtls_ecp_group_id grp_id =
-            mbedtls_ecc_group_of_psa( curve, PSA_BITS_TO_BYTES( attributes->core.bits ) );
+            mbedtls_ecc_group_of_psa( curve, PSA_BITS_TO_BYTES( psa_get_key_bits( attributes ) ) );
         const mbedtls_ecp_curve_info *curve_info =
             mbedtls_ecp_curve_info_from_grp_id( grp_id );
         mbedtls_ecp_keypair ecp;
@@ -79,7 +81,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         if( grp_id == MBEDTLS_ECP_DP_NONE || curve_info == NULL )
             return( PSA_ERROR_NOT_SUPPORTED );
-        if( curve_info->bit_size != attributes->core.bits )
+        if( curve_info->bit_size != psa_get_key_bits( attributes ) )
             return( PSA_ERROR_INVALID_ARGUMENT );
         mbedtls_ecp_keypair_init( &ecp );
         ret = mbedtls_ecp_gen_key( grp_id, &ecp,
@@ -92,7 +94,7 @@
         }
 
         /* Make sure to use export representation */
-        size_t bytes = PSA_BITS_TO_BYTES( attributes->core.bits );
+        size_t bytes = PSA_BITS_TO_BYTES( psa_get_key_bits( attributes ) );
         if( key_size < bytes )
         {
             mbedtls_ecp_keypair_free( &ecp );
diff --git a/tests/src/drivers/signature.c b/tests/src/drivers/signature.c
index 04c5de4..d1a6009 100644
--- a/tests/src/drivers/signature.c
+++ b/tests/src/drivers/signature.c
@@ -1,5 +1,7 @@
 /*
- * Test driver for signature functions
+ * Test driver for signature functions.
+ * Currently supports signing and verifying precalculated hashes, using
+ * only deterministic ECDSA on curves secp256r1, secp384r1 and secp521r1.
  */
 /*  Copyright The Mbed TLS Contributors
  *  SPDX-License-Identifier: Apache-2.0
@@ -130,10 +132,7 @@
                                                signature + curve_bytes,
                                                curve_bytes ) );
 cleanup:
-    /* There's no easy way to translate the error code except through a
-     * library function that's not exported. Use a debugger. */
-    if( ret == 0 )
-        status = PSA_SUCCESS;
+    status = mbedtls_to_psa_error( ret );
     mbedtls_mpi_free( &r );
     mbedtls_mpi_free( &s );
     mbedtls_ecp_keypair_free( &ecp );
@@ -258,10 +257,7 @@
     MBEDTLS_MPI_CHK( mbedtls_ecdsa_verify( &ecp.grp, hash, hash_length,
                                 &ecp.Q, &r, &s ) );
 cleanup:
-    /* There's no easy way to translate the error code except through a
-     * library function that's not exported. Use a debugger. */
-    if( ret == 0 )
-        status = PSA_SUCCESS;
+    status = mbedtls_to_psa_error( ret );
     mbedtls_mpi_free( &r );
     mbedtls_mpi_free( &s );
     mbedtls_ecp_keypair_free( &ecp );