diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index c010083..381fc63 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -3482,6 +3482,14 @@
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
 aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:-1:"":"":PSA_ERROR_INVALID_ARGUMENT
 
+PSA AEAD output buffer test: AES - CCM, IN = 40 BUF = 39
+depends_on:PSA_WANT_ALG_CCM:PSA_WANT_KEY_TYPE_AES
+aead_multipart_update_buffer_test:PSA_KEY_TYPE_AES:"4189351B5CAEA375A0299E81C621BF43":PSA_ALG_CCM:39:"48c0906930561e0ab0ef4cd972":"40a27c1d1e23ea3dbe8056b2774861a4a201cce49f19997d19206d8c8a343951":"26d56961c035a7e452cce61bc6ee220d77b3f94d18fd10b6d80e8bf80f4a46cab06d4313f0db9be9":PSA_ERROR_BUFFER_TOO_SMALL
+
+PSA AEAD output buffer test: AES - CCM, IN = 40 BUF = 0
+depends_on:PSA_WANT_ALG_CCM:PSA_WANT_KEY_TYPE_AES
+aead_multipart_update_buffer_test:PSA_KEY_TYPE_AES:"4189351B5CAEA375A0299E81C621BF43":PSA_ALG_CCM:0:"48c0906930561e0ab0ef4cd972":"40a27c1d1e23ea3dbe8056b2774861a4a201cce49f19997d19206d8c8a343951":"26d56961c035a7e452cce61bc6ee220d77b3f94d18fd10b6d80e8bf80f4a46cab06d4313f0db9be9":PSA_ERROR_BUFFER_TOO_SMALL
+
 PSA AEAD output buffer test: AES - GCM, IN = 16, BUF = 15
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
 aead_multipart_update_buffer_test:PSA_KEY_TYPE_AES:"fbc0b4c56a714c83217b2d1bcadd2ed2e9efb0dcac6cc19f":PSA_ALG_GCM:15:"5f4b43e811da9c470d6a9b01":"":"d2ae38c4375954835d75b8e4c2f9bbb4":PSA_ERROR_BUFFER_TOO_SMALL
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 3e1355d..8b73ae7 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4238,6 +4238,9 @@
 
     PSA_ASSERT( status );
 
+    PSA_ASSERT( psa_aead_set_lengths( &operation, additional_data->len,
+                                      input_data->len ) );
+
     PSA_ASSERT( psa_aead_set_nonce( &operation, nonce->x, nonce->len ) );
 
     PSA_ASSERT( psa_aead_update_ad( &operation, additional_data->x,