diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 976b7822f..967d5e3eb 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4883,6 +4883,80 @@ void aead_multipart_state_test( int key_type_arg, data_t *key_data,
 
     psa_aead_abort( &operation );
 
+    /* ------------------------------------------------------- */ 
+    /* Test for setting nonce after calling set lengths */
+
+    PSA_ASSERT( psa_aead_encrypt_setup( &operation, key, alg ) );
+
+    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_aead_abort( &operation );
+
+    /* Test for setting nonce after calling set lengths with UINT32_MAX length */
+
+    PSA_ASSERT( psa_aead_encrypt_setup( &operation, key, alg ) );
+
+    if( operation.alg == PSA_ALG_CCM )
+    {
+        TEST_EQUAL( psa_aead_set_lengths( &operation, UINT32_MAX,
+                                          input_data->len ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        TEST_EQUAL( psa_aead_set_nonce( &operation, nonce->x, nonce->len ),
+                    PSA_ERROR_BAD_STATE );
+    }
+    else
+    {
+        PSA_ASSERT( psa_aead_set_lengths( &operation, UINT32_MAX,
+                                          input_data->len ) );
+        PSA_ASSERT( psa_aead_set_nonce( &operation, nonce->x, nonce->len ) );
+    }
+
+    psa_aead_abort( &operation );
+
+    /* Test for setting nonce after calling set lengths with SIZE_MAX length */
+#if SIZE_MAX > UINT32_MAX
+    PSA_ASSERT( psa_aead_encrypt_setup( &operation, key, alg ) );
+
+    if( operation.alg == PSA_ALG_CCM || operation.alg == PSA_ALG_GCM )
+    {
+        TEST_EQUAL( psa_aead_set_lengths( &operation, SIZE_MAX,
+                                          input_data->len ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+        TEST_EQUAL( psa_aead_set_nonce( &operation, nonce->x, nonce->len ),
+                    PSA_ERROR_BAD_STATE );
+    }
+    else
+    {
+        PSA_ASSERT( psa_aead_set_lengths( &operation, SIZE_MAX,
+                                          input_data->len ) );
+        PSA_ASSERT( psa_aead_set_nonce( &operation, nonce->x, nonce->len ) );
+    }
+
+    psa_aead_abort( &operation );
+#endif
+
+    /* Test for calling set lengths with a length too long, after setting nonce */
+
+    PSA_ASSERT( psa_aead_encrypt_setup( &operation, key, alg ) );
+
+    PSA_ASSERT( psa_aead_set_nonce( &operation, nonce->x, nonce->len ) );
+
+    if( operation.alg == PSA_ALG_CCM )
+    {
+        TEST_EQUAL( psa_aead_set_lengths( &operation, UINT32_MAX,
+                                          input_data->len ),
+                    PSA_ERROR_INVALID_ARGUMENT );
+    }
+    else
+    {
+        PSA_ASSERT( psa_aead_set_lengths( &operation, UINT32_MAX,
+                                          input_data->len ) );
+    }
+
+    psa_aead_abort( &operation );
 
     /* ------------------------------------------------------- */