diff --git a/ChangeLog.d/psa_raw_key_agreement-buffer_too_small.txt b/ChangeLog.d/psa_raw_key_agreement-buffer_too_small.txt
new file mode 100644
index 0000000..415c849
--- /dev/null
+++ b/ChangeLog.d/psa_raw_key_agreement-buffer_too_small.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * psa_raw_key_agreement() now returns PSA_ERROR_BUFFER_TOO_SMALL when
+     applicable. Fixes #5735.
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c74aa15..8df26f0 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5019,6 +5019,22 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
+    /* PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE() is in general an upper bound
+     * for the output size. The PSA specification only guarantees that this
+     * function works if output_size >= PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE(...),
+     * but it might be nice to allow smaller buffers if the output fits.
+     * At the time of writing this comment, with only ECDH implemented,
+     * PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE() is exact so the point is moot.
+     * If FFDH is implemented, PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE() can easily
+     * be exact for it as well. */
+    size_t expected_length =
+        PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE( slot->attr.type, slot->attr.bits );
+    if( output_size < expected_length )
+    {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
+
     status = psa_key_agreement_raw_internal( alg, slot,
                                              peer_key, peer_key_length,
                                              output, output_size,
diff --git a/tests/include/test/helpers.h b/tests/include/test/helpers.h
index 5d10190..540f163 100644
--- a/tests/include/test/helpers.h
+++ b/tests/include/test/helpers.h
@@ -154,6 +154,48 @@
                         unsigned long long value1, unsigned long long value2 );
 
 /**
+ * \brief           Record the current test case as a failure based
+ *                  on comparing two unsigned integers.
+ *
+ *                  This function is usually called via the macro
+ *                  #TEST_LE_U.
+ *
+ * \param test      Description of the failure or assertion that failed. This
+ *                  MUST be a string literal. This normally has the form
+ *                  "EXPR1 <= EXPR2" where EXPR1 has the value \p value1
+ *                  and EXPR2 has the value \p value2.
+ * \param line_no   Line number where the failure originated.
+ * \param filename  Filename where the failure originated.
+ * \param value1    The first value to compare.
+ * \param value2    The second value to compare.
+ *
+ * \return          \c 1 if \p value1 <= \p value2, otherwise \c 0.
+ */
+int mbedtls_test_le_u( const char *test, int line_no, const char* filename,
+                       unsigned long long value1, unsigned long long value2 );
+
+/**
+ * \brief           Record the current test case as a failure based
+ *                  on comparing two signed integers.
+ *
+ *                  This function is usually called via the macro
+ *                  #TEST_LE_S.
+ *
+ * \param test      Description of the failure or assertion that failed. This
+ *                  MUST be a string literal. This normally has the form
+ *                  "EXPR1 <= EXPR2" where EXPR1 has the value \p value1
+ *                  and EXPR2 has the value \p value2.
+ * \param line_no   Line number where the failure originated.
+ * \param filename  Filename where the failure originated.
+ * \param value1    The first value to compare.
+ * \param value2    The second value to compare.
+ *
+ * \return          \c 1 if \p value1 <= \p value2, otherwise \c 0.
+ */
+int mbedtls_test_le_s( const char *test, int line_no, const char* filename,
+                       long long value1, long long value2 );
+
+/**
  * \brief          This function decodes the hexadecimal representation of
  *                 data.
  *
diff --git a/tests/include/test/macros.h b/tests/include/test/macros.h
index 7b4c2f6..2d68f7d 100644
--- a/tests/include/test/macros.h
+++ b/tests/include/test/macros.h
@@ -100,6 +100,32 @@
             goto exit;                                                  \
     } while( 0 )
 
+/** Evaluate two unsigned integer expressions and fail the test case
+ * if they are not in increasing order (left <= right).
+ *
+ * \param expr1     An integral-typed expression to evaluate.
+ * \param expr2     Another integral-typed expression to evaluate.
+ */
+#define TEST_LE_U( expr1, expr2 )                                       \
+    do {                                                                \
+        if( ! mbedtls_test_le_u( #expr1 " <= " #expr2, __LINE__, __FILE__, \
+                                 expr1, expr2 ) )                      \
+            goto exit;                                                  \
+    } while( 0 )
+
+/** Evaluate two signed integer expressions and fail the test case
+ * if they are not in increasing order (left <= right).
+ *
+ * \param expr1     An integral-typed expression to evaluate.
+ * \param expr2     Another integral-typed expression to evaluate.
+ */
+#define TEST_LE_S( expr1, expr2 )                                       \
+    do {                                                                \
+        if( ! mbedtls_test_le_s( #expr1 " <= " #expr2, __LINE__, __FILE__, \
+                                 expr1, expr2 ) )                      \
+            goto exit;                                                  \
+    } while( 0 )
+
 /** Allocate memory dynamically and fail the test case if this fails.
  * The allocated memory will be filled with zeros.
  *
diff --git a/tests/scripts/test_psa_compliance.py b/tests/scripts/test_psa_compliance.py
index b3bfd31..777ffcb 100755
--- a/tests/scripts/test_psa_compliance.py
+++ b/tests/scripts/test_psa_compliance.py
@@ -50,7 +50,7 @@
 #
 # Web URL: https://github.com/bensze01/psa-arch-tests/tree/fixes-for-mbedtls-2
 PSA_ARCH_TESTS_REPO = 'https://github.com/bensze01/psa-arch-tests.git'
-PSA_ARCH_TESTS_REF = 'fixes-for-mbedtls-2'
+PSA_ARCH_TESTS_REF = 'fix-5735-2.28'
 
 #pylint: disable=too-many-branches,too-many-statements
 def main():
diff --git a/tests/src/helpers.c b/tests/src/helpers.c
index c53e4fc..abc6852 100644
--- a/tests/src/helpers.c
+++ b/tests/src/helpers.c
@@ -140,6 +140,52 @@
     return( 0 );
 }
 
+int mbedtls_test_le_u( const char *test, int line_no, const char* filename,
+                       unsigned long long value1, unsigned long long value2 )
+{
+    if( value1 <= value2 )
+        return( 1 );
+    if( mbedtls_test_info.result == MBEDTLS_TEST_RESULT_FAILED )
+    {
+        /* We've already recorded the test as having failed. Don't
+         * overwrite any previous information about the failure. */
+        return( 0 );
+    }
+    mbedtls_test_fail( test, line_no, filename );
+    (void) mbedtls_snprintf( mbedtls_test_info.line1,
+                             sizeof( mbedtls_test_info.line1 ),
+                             "lhs = 0x%016llx = %llu",
+                             value1, value1 );
+    (void) mbedtls_snprintf( mbedtls_test_info.line2,
+                             sizeof( mbedtls_test_info.line2 ),
+                             "rhs = 0x%016llx = %llu",
+                             value2, value2 );
+    return( 0 );
+}
+
+int mbedtls_test_le_s( const char *test, int line_no, const char* filename,
+                       long long value1, long long value2 )
+{
+    if( value1 <= value2 )
+        return( 1 );
+    if( mbedtls_test_info.result == MBEDTLS_TEST_RESULT_FAILED )
+    {
+        /* We've already recorded the test as having failed. Don't
+         * overwrite any previous information about the failure. */
+        return( 0 );
+    }
+    mbedtls_test_fail( test, line_no, filename );
+    (void) mbedtls_snprintf( mbedtls_test_info.line1,
+                             sizeof( mbedtls_test_info.line1 ),
+                             "lhs = 0x%016llx = %lld",
+                             (unsigned long long) value1, value1 );
+    (void) mbedtls_snprintf( mbedtls_test_info.line2,
+                             sizeof( mbedtls_test_info.line2 ),
+                             "rhs = 0x%016llx = %lld",
+                             (unsigned long long) value2, value2 );
+    return( 0 );
+}
+
 int mbedtls_test_unhexify( unsigned char *obuf,
                            size_t obufmax,
                            const char *ibuf,
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 8950f8d..f9e9093 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -289,7 +289,7 @@
     /* Check that the length for a truncated MAC always fits in the algorithm
      * encoding. The shifted mask is the maximum truncated value. The
      * untruncated algorithm may be one byte larger. */
-    TEST_ASSERT( PSA_MAC_MAX_SIZE <= 1 + max_truncated_mac_size );
+    TEST_LE_U( PSA_MAC_MAX_SIZE, 1 + max_truncated_mac_size );
 
 #if defined(MBEDTLS_TEST_DEPRECATED)
     /* Check deprecated constants. */
@@ -591,7 +591,7 @@
      * and export_size. On errors, the exported length must be 0. */
     TEST_ASSERT( exported_length != INVALID_EXPORT_LENGTH );
     TEST_ASSERT( status == PSA_SUCCESS || exported_length == 0 );
-    TEST_ASSERT( exported_length <= export_size );
+    TEST_LE_U( exported_length, export_size );
 
     TEST_ASSERT( mem_is_char( exported + exported_length, 0,
                               export_size - exported_length ) );
@@ -626,7 +626,7 @@
     TEST_ASSERT( exported_length <=
                  PSA_EXPORT_KEY_OUTPUT_SIZE( type,
                                              psa_get_key_bits( &got_attributes ) ) );
-    TEST_ASSERT( exported_length <= PSA_EXPORT_KEY_PAIR_MAX_SIZE );
+    TEST_LE_U( exported_length, PSA_EXPORT_KEY_PAIR_MAX_SIZE );
 
 destroy:
     /* Destroy the key */
@@ -685,12 +685,12 @@
         size_t bits;
         PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
         bits = psa_get_key_bits( &attributes );
-        TEST_ASSERT( expected_public_key->len <=
-                     PSA_EXPORT_KEY_OUTPUT_SIZE( public_type, bits ) );
-        TEST_ASSERT( expected_public_key->len <=
-                     PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( public_type, bits ) );
-        TEST_ASSERT( expected_public_key->len <=
-                     PSA_EXPORT_PUBLIC_KEY_MAX_SIZE );
+        TEST_LE_U( expected_public_key->len,
+                   PSA_EXPORT_KEY_OUTPUT_SIZE( public_type, bits ) );
+        TEST_LE_U( expected_public_key->len,
+                   PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE( public_type, bits ) );
+        TEST_LE_U( expected_public_key->len,
+                   PSA_EXPORT_PUBLIC_KEY_MAX_SIZE );
         ASSERT_COMPARE( expected_public_key->x, expected_public_key->len,
                         exported, exported_length );
     }
@@ -983,8 +983,8 @@
     size_t tag_length = tag_length_arg;
     size_t output_length;
 
-    TEST_ASSERT( nonce_length <= sizeof( nonce ) );
-    TEST_ASSERT( tag_length <= sizeof( tag ) );
+    TEST_LE_U( nonce_length, sizeof( nonce ) );
+    TEST_LE_U( tag_length, sizeof( tag ) );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -1576,7 +1576,7 @@
     status = psa_hash_compute( alg, input->x, input->len,
                                output, output_size, &output_length );
     TEST_EQUAL( status, expected_status );
-    TEST_ASSERT( output_length <= output_size );
+    TEST_LE_U( output_length, output_size );
 
 exit:
     mbedtls_free( output );
@@ -2132,7 +2132,7 @@
         expected_mac->len + 1,
     };
 
-    TEST_ASSERT( mac_buffer_size <= PSA_MAC_MAX_SIZE );
+    TEST_LE_U( mac_buffer_size, PSA_MAC_MAX_SIZE );
     /* We expect PSA_MAC_LENGTH to be exact. */
     TEST_ASSERT( expected_mac->len == mac_buffer_size );
 
@@ -2210,7 +2210,7 @@
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     uint8_t *perturbed_mac = NULL;
 
-    TEST_ASSERT( expected_mac->len <= PSA_MAC_MAX_SIZE );
+    TEST_LE_U( expected_mac->len, PSA_MAC_MAX_SIZE );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -2608,14 +2608,14 @@
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Validate size macros */
-    TEST_ASSERT( ciphertext->len <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, plaintext->len ) );
-    TEST_ASSERT( PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, plaintext->len ) <=
+    TEST_LE_U( ciphertext->len,
+               PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, plaintext->len ) );
+    TEST_LE_U( PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, plaintext->len ),
                  PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( plaintext->len ) );
-    TEST_ASSERT( plaintext->len <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, ciphertext->len ) );
-    TEST_ASSERT( PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, ciphertext->len ) <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( ciphertext->len ) );
+    TEST_LE_U( plaintext->len,
+               PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, ciphertext->len ) );
+    TEST_LE_U( PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, ciphertext->len ),
+               PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( ciphertext->len ) );
 
 
     /* Set up key and output buffer */
@@ -2655,7 +2655,7 @@
                                    plaintext->x, plaintext->len,
                                    output, output_buffer_size,
                                    &length ) );
-    TEST_ASSERT( length <= output_buffer_size );
+    TEST_LE_U( length, output_buffer_size );
     output_length += length;
     PSA_ASSERT( psa_cipher_finish( &operation,
                                    output + output_length,
@@ -2673,7 +2673,7 @@
                                    ciphertext->x, ciphertext->len,
                                    output, output_buffer_size,
                                    &length ) );
-    TEST_ASSERT( length <= output_buffer_size );
+    TEST_LE_U( length, output_buffer_size );
     output_length += length;
     PSA_ASSERT( psa_cipher_finish( &operation,
                                    output + output_length,
@@ -2786,10 +2786,10 @@
        the output is not possible. Validating with multipart encryption. */
     PSA_ASSERT( psa_cipher_encrypt( key, alg, input->x, input->len, output1,
                                     output1_buffer_size, &output1_length ) );
-    TEST_ASSERT( output1_length <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len ) );
-    TEST_ASSERT( output1_length <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
+    TEST_LE_U( output1_length,
+               PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len ) );
+    TEST_LE_U( output1_length,
+               PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
 
     PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
     PSA_ASSERT( psa_cipher_set_iv( &operation, output1, iv_size ) );
@@ -2798,20 +2798,20 @@
                                    input->x, input->len,
                                    output2, output2_buffer_size,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, input->len ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, input->len ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
     output2_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_finish( &operation,
                                    output2 + output2_length,
                                    output2_buffer_size - output2_length,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
     output2_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_abort( &operation ) );
@@ -2871,15 +2871,15 @@
                          PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg );
     ASSERT_ALLOC( output, output_buffer_size );
 
-    TEST_ASSERT( first_part_size <= input->len );
+    TEST_LE_U( first_part_size, input->len );
     PSA_ASSERT( psa_cipher_update( &operation, input->x, first_part_size,
                                    output, output_buffer_size,
                                    &function_output_length ) );
     TEST_ASSERT( function_output_length == output1_length );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size) );
     total_output_length += function_output_length;
 
     if( first_part_size < input->len )
@@ -2892,12 +2892,12 @@
                                        output_buffer_size - total_output_length,
                                        &function_output_length ) );
         TEST_ASSERT( function_output_length == output2_length );
-        TEST_ASSERT( function_output_length <=
-                     PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
-                                                    alg,
-                                                    input->len - first_part_size ) );
-        TEST_ASSERT( function_output_length <=
-                     PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
+        TEST_LE_U( function_output_length,
+                   PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
+                                                  alg,
+                                                  input->len - first_part_size ) );
+        TEST_LE_U( function_output_length,
+                   PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
         total_output_length += function_output_length;
     }
 
@@ -2906,10 +2906,10 @@
                                   output + total_output_length ),
                                 output_buffer_size - total_output_length,
                                 &function_output_length );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
     total_output_length += function_output_length;
     TEST_EQUAL( status, expected_status );
 
@@ -2973,16 +2973,16 @@
                          PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg );
     ASSERT_ALLOC( output, output_buffer_size );
 
-    TEST_ASSERT( first_part_size <= input->len );
+    TEST_LE_U( first_part_size, input->len );
     PSA_ASSERT( psa_cipher_update( &operation,
                                    input->x, first_part_size,
                                    output, output_buffer_size,
                                    &function_output_length ) );
     TEST_ASSERT( function_output_length == output1_length );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
     total_output_length += function_output_length;
 
     if( first_part_size < input->len )
@@ -2995,12 +2995,12 @@
                                        output_buffer_size - total_output_length,
                                        &function_output_length ) );
         TEST_ASSERT( function_output_length == output2_length );
-        TEST_ASSERT( function_output_length <=
-                     PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
-                                                    alg,
-                                                    input->len - first_part_size ) );
-        TEST_ASSERT( function_output_length <=
-                     PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
+        TEST_LE_U( function_output_length,
+                   PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
+                                                  alg,
+                                                  input->len - first_part_size ) );
+        TEST_LE_U( function_output_length,
+                   PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len ) );
         total_output_length += function_output_length;
     }
 
@@ -3009,10 +3009,10 @@
                                   output + total_output_length ),
                                 output_buffer_size - total_output_length,
                                 &function_output_length );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
     total_output_length += function_output_length;
     TEST_EQUAL( status, expected_status );
 
@@ -3129,10 +3129,10 @@
 
     PSA_ASSERT( psa_cipher_decrypt( key, alg, input, input_buffer_size, output,
                                     output_buffer_size, &output_length ) );
-    TEST_ASSERT( output_length <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, input_buffer_size ) );
-    TEST_ASSERT( output_length <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( input_buffer_size ) );
+    TEST_LE_U( output_length,
+               PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, input_buffer_size ) );
+    TEST_LE_U( output_length,
+               PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( input_buffer_size ) );
 
     ASSERT_COMPARE( expected_output->x, expected_output->len,
                     output, output_length );
@@ -3175,10 +3175,10 @@
     PSA_ASSERT( psa_cipher_encrypt( key, alg, input->x, input->len,
                                     output1, output1_size,
                                     &output1_length ) );
-    TEST_ASSERT( output1_length <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len ) );
-    TEST_ASSERT( output1_length <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
+    TEST_LE_U( output1_length,
+               PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len ) );
+    TEST_LE_U( output1_length,
+               PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
 
     output2_size = output1_length;
     ASSERT_ALLOC( output2, output2_size );
@@ -3186,10 +3186,10 @@
     PSA_ASSERT( psa_cipher_decrypt( key, alg, output1, output1_length,
                                     output2, output2_size,
                                     &output2_length ) );
-    TEST_ASSERT( output2_length <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, output1_length ) );
-    TEST_ASSERT( output2_length <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( output1_length ) );
+    TEST_LE_U( output2_length,
+               PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, output1_length ) );
+    TEST_LE_U( output2_length,
+               PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( output1_length ) );
 
     ASSERT_COMPARE( input->x, input->len, output2, output2_length );
 
@@ -3246,19 +3246,19 @@
     }
 
     output1_buffer_size = PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len );
-    TEST_ASSERT( output1_buffer_size <=
-                 PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
+    TEST_LE_U( output1_buffer_size,
+               PSA_CIPHER_ENCRYPT_OUTPUT_MAX_SIZE( input->len ) );
     ASSERT_ALLOC( output1, output1_buffer_size );
 
-    TEST_ASSERT( first_part_size <= input->len );
+    TEST_LE_U( first_part_size, input->len );
 
     PSA_ASSERT( psa_cipher_update( &operation1, input->x, first_part_size,
                                    output1, output1_buffer_size,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
     output1_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_update( &operation1,
@@ -3266,31 +3266,31 @@
                                    input->len - first_part_size,
                                    output1, output1_buffer_size,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
-                                                alg,
-                                                input->len - first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len - first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
+                                              alg,
+                                              input->len - first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( input->len - first_part_size ) );
     output1_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_finish( &operation1,
                                    output1 + output1_length,
                                    output1_buffer_size - output1_length,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
     output1_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_abort( &operation1 ) );
 
     output2_buffer_size = output1_length;
-    TEST_ASSERT( output2_buffer_size <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, output1_length ) );
-    TEST_ASSERT( output2_buffer_size <=
-                 PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( output1_length ) );
+    TEST_LE_U( output2_buffer_size,
+               PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, output1_length ) );
+    TEST_LE_U( output2_buffer_size,
+               PSA_CIPHER_DECRYPT_OUTPUT_MAX_SIZE( output1_length ) );
     ASSERT_ALLOC( output2, output2_buffer_size );
 
     if( iv_length > 0 )
@@ -3302,10 +3302,10 @@
     PSA_ASSERT( psa_cipher_update( &operation2, output1, first_part_size,
                                    output2, output2_buffer_size,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( first_part_size ) );
     output2_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_update( &operation2,
@@ -3313,22 +3313,22 @@
                                    output1_length - first_part_size,
                                    output2, output2_buffer_size,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
-                                                alg,
-                                                output1_length - first_part_size ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( output1_length - first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type,
+                                              alg,
+                                              output1_length - first_part_size ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_UPDATE_OUTPUT_MAX_SIZE( output1_length - first_part_size ) );
     output2_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_finish( &operation2,
                                    output2 + output2_length,
                                    output2_buffer_size - output2_length,
                                    &function_output_length ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
-    TEST_ASSERT( function_output_length <=
-                 PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg ) );
+    TEST_LE_U( function_output_length,
+               PSA_CIPHER_FINISH_OUTPUT_MAX_SIZE );
     output2_length += function_output_length;
 
     PSA_ASSERT( psa_cipher_abort( &operation2 ) );
@@ -3635,7 +3635,7 @@
     signature_size = PSA_SIGN_OUTPUT_SIZE( key_type,
                                                       key_bits, alg );
     TEST_ASSERT( signature_size != 0 );
-    TEST_ASSERT( signature_size <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_size, PSA_SIGNATURE_MAX_SIZE );
     ASSERT_ALLOC( signature, signature_size );
 
     /* Perform the signature. */
@@ -3706,7 +3706,7 @@
      * whatever it is, it should be less than signature_size, so that
      * if the caller tries to read *signature_length bytes without
      * checking the error code then they don't overflow a buffer. */
-    TEST_ASSERT( signature_length <= signature_size );
+    TEST_LE_U( signature_length, signature_size );
 
 #if defined(MBEDTLS_TEST_DEPRECATED)
     signature_length = INVALID_EXPORT_LENGTH;
@@ -3715,7 +3715,7 @@
                                      signature, signature_size,
                                      &signature_length ),
                 expected_status );
-    TEST_ASSERT( signature_length <= signature_size );
+    TEST_LE_U( signature_length, signature_size );
 #endif /* MBEDTLS_TEST_DEPRECATED */
 
 exit:
@@ -3755,7 +3755,7 @@
     signature_size = PSA_SIGN_OUTPUT_SIZE( key_type,
                                                       key_bits, alg );
     TEST_ASSERT( signature_size != 0 );
-    TEST_ASSERT( signature_size <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_size, PSA_SIGNATURE_MAX_SIZE );
     ASSERT_ALLOC( signature, signature_size );
 
     /* Perform the signature. */
@@ -3764,7 +3764,7 @@
                                signature, signature_size,
                                &signature_length ) );
     /* Check that the signature length looks sensible. */
-    TEST_ASSERT( signature_length <= signature_size );
+    TEST_LE_U( signature_length, signature_size );
     TEST_ASSERT( signature_length > 0 );
 
     /* Use the library to verify that the signature is correct. */
@@ -3807,7 +3807,7 @@
     psa_algorithm_t alg = alg_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    TEST_ASSERT( signature_data->len <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_data->len, PSA_SIGNATURE_MAX_SIZE );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -3907,7 +3907,7 @@
 
     signature_size = PSA_SIGN_OUTPUT_SIZE( key_type, key_bits, alg );
     TEST_ASSERT( signature_size != 0 );
-    TEST_ASSERT( signature_size <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_size, PSA_SIGNATURE_MAX_SIZE );
     ASSERT_ALLOC( signature, signature_size );
 
     PSA_ASSERT( psa_sign_message( key, alg,
@@ -3966,7 +3966,7 @@
      * whatever it is, it should be less than signature_size, so that
      * if the caller tries to read *signature_length bytes without
      * checking the error code then they don't overflow a buffer. */
-    TEST_ASSERT( signature_length <= signature_size );
+    TEST_LE_U( signature_length, signature_size );
 
 exit:
     psa_reset_key_attributes( &attributes );
@@ -4005,14 +4005,14 @@
 
     signature_size = PSA_SIGN_OUTPUT_SIZE( key_type, key_bits, alg );
     TEST_ASSERT( signature_size != 0 );
-    TEST_ASSERT( signature_size <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_size, PSA_SIGNATURE_MAX_SIZE );
     ASSERT_ALLOC( signature, signature_size );
 
     PSA_ASSERT( psa_sign_message( key, alg,
                                   input_data->x, input_data->len,
                                   signature, signature_size,
                                   &signature_length ) );
-    TEST_ASSERT( signature_length <= signature_size );
+    TEST_LE_U( signature_length, signature_size );
     TEST_ASSERT( signature_length > 0 );
 
     PSA_ASSERT( psa_verify_message( key, alg,
@@ -4052,7 +4052,7 @@
     psa_algorithm_t alg = alg_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    TEST_ASSERT( signature_data->len <= PSA_SIGNATURE_MAX_SIZE );
+    TEST_LE_U( signature_data->len, PSA_SIGNATURE_MAX_SIZE );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -4146,7 +4146,7 @@
     key_bits = psa_get_key_bits( &attributes );
 
     output_size = PSA_ASYMMETRIC_ENCRYPT_OUTPUT_SIZE( key_type, key_bits, alg );
-    TEST_ASSERT( output_size <= PSA_ASYMMETRIC_ENCRYPT_OUTPUT_MAX_SIZE );
+    TEST_LE_U( output_size, PSA_ASYMMETRIC_ENCRYPT_OUTPUT_MAX_SIZE );
     ASSERT_ALLOC( output, output_size );
 
     /* Encrypt the input */
@@ -4220,13 +4220,13 @@
     key_bits = psa_get_key_bits( &attributes );
 
     output_size = PSA_ASYMMETRIC_ENCRYPT_OUTPUT_SIZE( key_type, key_bits, alg );
-    TEST_ASSERT( output_size <= PSA_ASYMMETRIC_ENCRYPT_OUTPUT_MAX_SIZE );
+    TEST_LE_U( output_size, PSA_ASYMMETRIC_ENCRYPT_OUTPUT_MAX_SIZE );
     ASSERT_ALLOC( output, output_size );
 
     output2_size = input_data->len;
-    TEST_ASSERT( output2_size <=
-                 PSA_ASYMMETRIC_DECRYPT_OUTPUT_SIZE( key_type, key_bits, alg ) );
-    TEST_ASSERT( output2_size <= PSA_ASYMMETRIC_DECRYPT_OUTPUT_MAX_SIZE );
+    TEST_LE_U( output2_size,
+               PSA_ASYMMETRIC_DECRYPT_OUTPUT_SIZE( key_type, key_bits, alg ) );
+    TEST_LE_U( output2_size, PSA_ASYMMETRIC_DECRYPT_OUTPUT_MAX_SIZE );
     ASSERT_ALLOC( output2, output2_size );
 
     /* We test encryption by checking that encrypt-then-decrypt gives back
@@ -4239,7 +4239,7 @@
                                         &output_length ) );
     /* We don't know what ciphertext length to expect, but check that
      * it looks sensible. */
-    TEST_ASSERT( output_length <= output_size );
+    TEST_LE_U( output_length, output_size );
 
     PSA_ASSERT( psa_asymmetric_decrypt( key, alg,
                                         output, output_length,
@@ -4294,7 +4294,7 @@
 
     /* Determine the maximum ciphertext length */
     output_size = PSA_ASYMMETRIC_DECRYPT_OUTPUT_SIZE( key_type, key_bits, alg );
-    TEST_ASSERT( output_size <= PSA_ASYMMETRIC_DECRYPT_OUTPUT_MAX_SIZE );
+    TEST_LE_U( output_size, PSA_ASYMMETRIC_DECRYPT_OUTPUT_MAX_SIZE );
     ASSERT_ALLOC( output, output_size );
 
     PSA_ASSERT( psa_asymmetric_decrypt( key, alg,
@@ -4367,7 +4367,7 @@
                                             output, output_size,
                                             &output_length );
     TEST_EQUAL( actual_status, expected_status );
-    TEST_ASSERT( output_length <= output_size );
+    TEST_LE_U( output_length, output_size );
 
     /* If the label is empty, the test framework puts a non-null pointer
      * in label->x. Test that a null pointer works as well. */
@@ -4382,7 +4382,7 @@
                                                 output, output_size,
                                                 &output_length );
         TEST_EQUAL( actual_status, expected_status );
-        TEST_ASSERT( output_length <= output_size );
+        TEST_LE_U( output_length, output_size );
     }
 
 exit:
@@ -5094,7 +5094,6 @@
     size_t output_length = ~0;
     size_t key_bits;
 
-    ASSERT_ALLOC( output, expected_output->len );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
@@ -5107,16 +5106,47 @@
     PSA_ASSERT( psa_get_key_attributes( our_key, &attributes ) );
     key_bits = psa_get_key_bits( &attributes );
 
+    /* Validate size macros */
+    TEST_LE_U( expected_output->len,
+               PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE( our_key_type, key_bits ) );
+    TEST_LE_U( PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE( our_key_type, key_bits ),
+               PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE );
+
+    /* Good case with exact output size */
+    ASSERT_ALLOC( output, expected_output->len );
     PSA_ASSERT( psa_raw_key_agreement( alg, our_key,
                                        peer_key_data->x, peer_key_data->len,
                                        output, expected_output->len,
                                        &output_length ) );
     ASSERT_COMPARE( output, output_length,
                     expected_output->x, expected_output->len );
-    TEST_ASSERT( output_length <=
-                 PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE( our_key_type, key_bits ) );
-    TEST_ASSERT( output_length <=
-                 PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE );
+    mbedtls_free( output );
+    output = NULL;
+    output_length = ~0;
+
+    /* Larger buffer */
+    ASSERT_ALLOC( output, expected_output->len + 1 );
+    PSA_ASSERT( psa_raw_key_agreement( alg, our_key,
+                                       peer_key_data->x, peer_key_data->len,
+                                       output, expected_output->len + 1,
+                                       &output_length ) );
+    ASSERT_COMPARE( output, output_length,
+                    expected_output->x, expected_output->len );
+    mbedtls_free( output );
+    output = NULL;
+    output_length = ~0;
+
+    /* Buffer too small */
+    ASSERT_ALLOC( output, expected_output->len - 1 );
+    TEST_EQUAL( psa_raw_key_agreement( alg, our_key,
+                                       peer_key_data->x, peer_key_data->len,
+                                       output, expected_output->len - 1,
+                                       &output_length ),
+                PSA_ERROR_BUFFER_TOO_SMALL );
+    /* Not required by the spec, but good robustness */
+    TEST_LE_U( output_length, expected_output->len - 1 );
+    mbedtls_free( output );
+    output = NULL;
 
 exit:
     mbedtls_free( output );