diff --git a/tests/src/drivers/test_driver_aead.c b/tests/src/drivers/test_driver_aead.c
index d27ada2..6befe7c 100644
--- a/tests/src/drivers/test_driver_aead.c
+++ b/tests/src/drivers/test_driver_aead.c
@@ -171,8 +171,9 @@
     }
     else
     {
-       /* No mbedtls_psa_aead_set_lengths, everything is done in PSA Core. */
-       mbedtls_test_driver_aead_hooks.driver_status = PSA_SUCCESS;
+        mbedtls_test_driver_aead_hooks.driver_status =
+            mbedtls_psa_aead_set_lengths( operation, ad_length,
+                                          plaintext_length );
     }
 
     return( mbedtls_test_driver_aead_hooks.driver_status );