diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 9e17e27f2d..8fbfa77916 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5523,6 +5523,10 @@ exit:
 
 static psa_status_t psa_aead_final_checks(const psa_aead_operation_t *operation)
 {
+    if (operation->alg == PSA_ALG_CCM && !operation->lengths_set) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->id == 0 || !operation->nonce_set) {
         return PSA_ERROR_BAD_STATE;
     }
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 2c9a95d779..df3324d883 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -6496,7 +6496,7 @@ void aead_multipart_state_test(int key_type_arg, data_t *key_data,
                                    &output_part_length,
                                    tag_buffer, tag_length,
                                    &tag_size),
-                   PSA_ERROR_INVALID_ARGUMENT);
+                   PSA_ERROR_BAD_STATE);
         psa_aead_abort(&operation);
 
         PSA_ASSERT(psa_aead_encrypt_setup(&operation, key, alg));
@@ -6617,7 +6617,7 @@ void aead_multipart_state_test(int key_type_arg, data_t *key_data,
                                    &output_part_length,
                                    tag_buffer, tag_length,
                                    &tag_size),
-                   PSA_ERROR_INVALID_ARGUMENT);
+                   PSA_ERROR_BAD_STATE);
         psa_aead_abort(&operation);
 
         PSA_ASSERT(psa_aead_encrypt_setup(&operation, key, alg));