Test calling psa_aead_set_lengths and set_nonce in various order

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index cbac109..976b782 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4117,6 +4117,7 @@
 void aead_multipart_set_nonce( int key_type_arg, data_t *key_data,
                                int alg_arg,
                                int nonce_length_arg,
+                               int set_lengths_method_arg,
                                data_t *additional_data,
                                data_t *input_data,
                                int expected_status_arg )
@@ -4139,6 +4140,7 @@
     size_t tag_length = 0;
     uint8_t tag_buffer[PSA_AEAD_TAG_MAX_SIZE];
     size_t index = 0;
+    set_lengths_method_t set_lengths_method = set_lengths_method_arg;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -4196,26 +4198,40 @@
         }
     }
 
+    if( set_lengths_method == SET_LENGTHS_BEFORE_NONCE )
+    {
+        PSA_ASSERT( psa_aead_set_lengths( &operation, additional_data->len,
+                                          input_data->len ) );
+    }
+
     status = psa_aead_set_nonce( &operation, nonce_buffer, nonce_length );
 
     TEST_EQUAL( status, expected_status );
 
     if( expected_status == PSA_SUCCESS )
     {
-        /* Ensure we can still complete operation. */
-        PSA_ASSERT( psa_aead_set_lengths( &operation, additional_data->len,
-                                          input_data->len ) );
+        if( set_lengths_method == SET_LENGTHS_AFTER_NONCE )
+        {
+            PSA_ASSERT( psa_aead_set_lengths( &operation, additional_data->len,
+                                              input_data->len ) );
+        }
+        if( operation.alg == PSA_ALG_CCM && set_lengths_method == DO_NOT_SET_LENGTHS )
+            expected_status = PSA_ERROR_BAD_STATE;
 
-        PSA_ASSERT( psa_aead_update_ad( &operation, additional_data->x,
-                                        additional_data->len ) );
+        /* Ensure we can still complete operation, unless it's CCM and we didn't set lengths. */
+        TEST_EQUAL( psa_aead_update_ad( &operation, additional_data->x,
+                                        additional_data->len ),
+                    expected_status );
 
-        PSA_ASSERT( psa_aead_update( &operation, input_data->x, input_data->len,
+        TEST_EQUAL( psa_aead_update( &operation, input_data->x, input_data->len,
                                      output, output_size,
-                                     &ciphertext_length ) );
+                                     &ciphertext_length ),
+                    expected_status );
 
-        PSA_ASSERT( psa_aead_finish( &operation, ciphertext, ciphertext_size,
+        TEST_EQUAL( psa_aead_finish( &operation, ciphertext, ciphertext_size,
                                      &ciphertext_length, tag_buffer,
-                                     PSA_AEAD_TAG_MAX_SIZE, &tag_length ) );
+                                     PSA_AEAD_TAG_MAX_SIZE, &tag_length ),
+                    expected_status );
     }
 
 exit: