diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 9f80d7b..408227d 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -467,7 +467,7 @@
     private_key_type = psa_get_key_type( &attributes );
     key_bits = psa_get_key_bits( &attributes );
     public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
-    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
+    public_key_length = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( public_key_type, key_bits );
     ASSERT_ALLOC( public_key, public_key_length );
     PSA_ASSERT( psa_export_public_key( key, public_key, public_key_length,
                                        &public_key_length ) );
@@ -509,7 +509,7 @@
     private_key_type = psa_get_key_type( &attributes );
     key_bits = psa_get_key_bits( &attributes );
     public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
-    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
+    public_key_length = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( public_key_type, key_bits );
     ASSERT_ALLOC( public_key, public_key_length );
     PSA_ASSERT( psa_export_public_key( key,
                                        public_key, public_key_length,
@@ -518,6 +518,15 @@
     status = psa_raw_key_agreement( alg, key,
                                     public_key, public_key_length,
                                     output, sizeof( output ), &output_length );
+    if ( status == PSA_SUCCESS )
+    {
+        TEST_ASSERT( output_length <=
+                     PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE( private_key_type,
+                                                        key_bits ) );
+        TEST_ASSERT( output_length <=
+                     PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE );
+    }
+
 exit:
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
@@ -625,6 +634,8 @@
         if( ! mbedtls_test_asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
             goto exit;
         TEST_EQUAL( p, end );
+
+        TEST_ASSERT( exported_length <= PSA_EXPORT_KEY_PAIR_MAX_SIZE );
     }
     else
 #endif /* MBEDTLS_RSA_C */
@@ -634,6 +645,8 @@
     {
         /* Just the secret value */
         TEST_EQUAL( exported_length, PSA_BITS_TO_BYTES( bits ) );
+
+        TEST_ASSERT( exported_length <= PSA_EXPORT_KEY_PAIR_MAX_SIZE );
     }
     else
 #endif /* MBEDTLS_ECP_C */
@@ -658,6 +671,12 @@
         if( ! mbedtls_test_asn1_skip_integer( &p, end, 2, bits, 1 ) )
             goto exit;
         TEST_EQUAL( p, end );
+
+
+        TEST_ASSERT( exported_length <=
+                     PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( type, bits ) );
+        TEST_ASSERT( exported_length <=
+                     PSA_EXPORT_PUBLIC_KEY_MAX_SIZE );
     }
     else
 #endif /* MBEDTLS_RSA_C */
@@ -665,6 +684,12 @@
 #if defined(MBEDTLS_ECP_C)
     if( PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( type ) )
     {
+
+        TEST_ASSERT( exported_length <=
+                     PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( type, bits ) );
+        TEST_ASSERT( exported_length <=
+                     PSA_EXPORT_PUBLIC_KEY_MAX_SIZE );
+
         if( PSA_KEY_TYPE_ECC_GET_FAMILY( type ) == PSA_ECC_FAMILY_MONTGOMERY )
         {
             /* The representation of an ECC Montgomery public key is
@@ -785,8 +810,8 @@
 
     public_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(
         psa_get_key_type( &attributes ) );
-    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE( public_type,
-                                                psa_get_key_bits( &attributes ) );
+    exported_size = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( public_type,
+                                                       psa_get_key_bits( &attributes ) );
     ASSERT_ALLOC( exported, exported_size );
 
     PSA_ASSERT( psa_export_public_key( key,