test_driver_cipher: add forced return status for encrypt and set_iv

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/tests/include/test/drivers/cipher.h b/tests/include/test/drivers/cipher.h
index 67047af..2fe47e4 100644
--- a/tests/include/test/drivers/cipher.h
+++ b/tests/include/test/drivers/cipher.h
@@ -23,13 +23,17 @@
     /* If not PSA_SUCCESS, return this error code instead of processing the
      * function call. */
     psa_status_t forced_status;
+    psa_status_t forced_status_encrypt;
+    psa_status_t forced_status_set_iv;
     /* Count the amount of times one of the cipher driver functions is called. */
     unsigned long hits;
-    unsigned long cipher_encrypt_hits;
-    unsigned long cipher_update_forced_status_hits;
+    unsigned long hits_encrypt;
+    unsigned long hits_set_iv;
 } mbedtls_test_driver_cipher_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, PSA_SUCCESS, 0, 0, 0 }
+#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, \
+                                          PSA_SUCCESS, PSA_SUCCESS, PSA_SUCCESS, \
+                                          0, 0, 0 }
 static inline mbedtls_test_driver_cipher_hooks_t
 mbedtls_test_driver_cipher_hooks_init(void)
 {
diff --git a/tests/src/drivers/test_driver_cipher.c b/tests/src/drivers/test_driver_cipher.c
index b9da5d2..2bc751a 100644
--- a/tests/src/drivers/test_driver_cipher.c
+++ b/tests/src/drivers/test_driver_cipher.c
@@ -41,7 +41,7 @@
     size_t *output_length)
 {
     mbedtls_test_driver_cipher_hooks.hits++;
-    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits++;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt++;
 
     if (mbedtls_test_driver_cipher_hooks.forced_output != NULL) {
         if (output_size < mbedtls_test_driver_cipher_hooks.forced_output_length) {
@@ -59,6 +59,9 @@
     if (mbedtls_test_driver_cipher_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_cipher_hooks.forced_status;
     }
+    if (mbedtls_test_driver_cipher_hooks.forced_status_encrypt != PSA_SUCCESS) {
+        return mbedtls_test_driver_cipher_hooks.forced_status_encrypt;
+    }
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
     defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_CIPHER)
@@ -209,10 +212,14 @@
     size_t iv_length)
 {
     mbedtls_test_driver_cipher_hooks.hits++;
+    mbedtls_test_driver_cipher_hooks.hits_set_iv++;
 
     if (mbedtls_test_driver_cipher_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_cipher_hooks.forced_status;
     }
+    if (mbedtls_test_driver_cipher_hooks.forced_status_set_iv != PSA_SUCCESS) {
+        return mbedtls_test_driver_cipher_hooks.forced_status_set_iv;
+    }
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
     defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_CIPHER)
@@ -249,7 +256,6 @@
     }
 
     if (mbedtls_test_driver_cipher_hooks.forced_status != PSA_SUCCESS) {
-        ++mbedtls_test_driver_cipher_hooks.cipher_update_forced_status_hits;
         return mbedtls_test_driver_cipher_hooks.forced_status;
     }
 
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 309b395..ff449ac 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -1061,10 +1061,10 @@
                               &key));
 
     mbedtls_test_driver_cipher_hooks.hits = 0;
-    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     PSA_ASSERT(psa_cipher_encrypt(key, alg, input->x, input->len, output1,
                                   output1_buffer_size, &output1_length));
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
     PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
@@ -1475,34 +1475,25 @@
      * successful, then force driver error.
      */
     mbedtls_test_driver_cipher_hooks.hits = 0;
-    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     TEST_EQUAL(status, PSA_SUCCESS);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
-    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
+    mbedtls_test_driver_cipher_hooks.forced_status_encrypt = PSA_ERROR_GENERIC_ERROR;
     /* Set the output buffer in a given state. */
     for (size_t i = 0; i < output_buffer_size; i++) {
         output[i] = 0xa5;
     }
 
-    mbedtls_test_driver_cipher_hooks.cipher_update_forced_status_hits = 0;
-    mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-#if defined(MBEDTLS_AES_C)
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_encrypt_hits, 1);
-#else
-    /* The call to psa_cipher_encrypt() is intentionally supposed to fail on the
-     * 1st access to the driver since "forced_status" is set. However this
-     * initial access happens in psa_cipher_update() (random number generation
-     * for IV) so psa_cipher_encrypt() never gets called. */
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.cipher_update_forced_status_hits, 1);
-#endif
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     TEST_EQUAL(status, PSA_ERROR_GENERIC_ERROR);
     /*
      * Check that the output buffer is still in the same state.
@@ -1563,25 +1554,18 @@
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status);
     mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_set_iv = 0;
 
-    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
+    mbedtls_test_driver_cipher_hooks.forced_status_set_iv = PSA_ERROR_GENERIC_ERROR;
     /* Set the output buffer in a given state. */
     for (size_t i = 0; i < 16; i++) {
         output[i] = 0xa5;
     }
 
     status = psa_cipher_generate_iv(&operation, output, 16, &function_output_length);
-    /* When generating the IV fails, it should call abort too */
-#if defined(MBEDTLS_AES_C)
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 2);
-#else
-    /* Previously failed call to psa_cipher_encrypt() above caused PSA to abort
-     * the cipher operation related to RNG. Therefore this call to
-     * psa_cipher_generate_iv() will failed due to unitialized RNG. Only the
-     * last driver call to psa_cipher_abort() remains. */
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
-#endif
-    TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_set_iv, 1);
+    TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status_set_iv);
+    mbedtls_test_driver_cipher_hooks.forced_status_set_iv = PSA_SUCCESS;
     /*
      * Check that the output buffer is still in the same state.
      * This will fail if the output buffer is used by the core to pass the IV