Add test for one-shot cipher driver wrapper functions

Signed-off-by: gabor-mezei-arm <gabor.mezei@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index dbc5b1f..ac91278 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -117,29 +117,69 @@
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:PSA_WANT_ECC_SECP_R1_256
 export_key:PSA_ERROR_GENERIC_ERROR:"":PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):"":PSA_ERROR_GENERIC_ERROR
 
-PSA symmetric encrypt: AES-CTR, 16 bytes, good
+PSA symmetric encrypt sanity check: AES-CTR, 16 bytes, good
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"8f9408fe80a81d3e813da3c7b0b2bd32":0:PSA_SUCCESS:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e117393172a":"":0:PSA_SUCCESS:PSA_SUCCESS
 
-PSA symmetric encrypt: AES-CTR, 15 bytes, good
+PSA symmetric encrypt sanity check: AES-CTR, 15 bytes, good
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"8f9408fe80a81d3e813da3c7b0b2bd":0:PSA_SUCCESS:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e11739317":"":0:PSA_SUCCESS:PSA_SUCCESS
 
-PSA symmetric encrypt: AES-CTR, 16 bytes, fallback
+PSA symmetric encrypt sanity check: AES-CTR, 16 bytes, fallback
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"8f9408fe80a81d3e813da3c7b0b2bd32":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e117393172a":"":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
 
-PSA symmetric encrypt: AES-CTR, 15 bytes, fallback
+PSA symmetric encrypt sanity check: AES-CTR, 15 bytes, fallback
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"8f9408fe80a81d3e813da3c7b0b2bd":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e11739317":"":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
 
-PSA symmetric encrypt: AES-CTR, 16 bytes, fake
+PSA symmetric encrypt sanity check: AES-CTR, 16 bytes, fake
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e117393172a":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
 
-PSA symmetric encrypt: AES-CTR, 15 bytes, fake
+PSA symmetric encrypt sanity check: AES-CTR, 15 bytes, fake
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
-cipher_encrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
+cipher_encrypt_sanity:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e11739317":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
+
+PSA symmetric encrypt validation: AES-CTR, 16 bytes, good
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_validation:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e117393172a"
+
+PSA symmetric encrypt validation: AES-CTR, 15 bytes, good
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_validation:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e11739317"
+
+PSA symmetric encrypt validation: AES-CTR, 16 bytes, fallback
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
+cipher_encrypt_validation:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e117393172a"
+
+PSA symmetric encrypt validation: AES-CTR, 15 bytes, fallback
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
+cipher_encrypt_validation:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"6bc1bee22e409f96e93d7e11739317"
+
+PSA symmetric encrypt multipart: AES-CTR, 16 bytes, good
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"8f9408fe80a81d3e813da3c7b0b2bd32":0:PSA_SUCCESS:PSA_SUCCESS
+
+PSA symmetric encrypt multipart: AES-CTR, 15 bytes, good
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"8f9408fe80a81d3e813da3c7b0b2bd":0:PSA_SUCCESS:PSA_SUCCESS
+
+PSA symmetric encrypt multipart: AES-CTR, 16 bytes, fallback
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"8f9408fe80a81d3e813da3c7b0b2bd32":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+
+PSA symmetric encrypt multipart: AES-CTR, 15 bytes, fallback
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"8f9408fe80a81d3e813da3c7b0b2bd":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+
+PSA symmetric encrypt multipart: AES-CTR, 16 bytes, fake
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
+
+PSA symmetric encrypt multipart: AES-CTR, 15 bytes, fake
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_encrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e11739317":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
 
 PSA symmetric decrypt: AES-CTR, 16 bytes, good
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
@@ -153,6 +193,18 @@
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
 cipher_decrypt:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"396ee84fb75fdbb5c2b13c7fe5a654aa":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
 
+PSA symmetric decrypt multipart: AES-CTR, 16 bytes, good
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_decrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"396ee84fb75fdbb5c2b13c7fe5a654aa":"dd3b5e5319b7591daab1e1a92687feb2":0:PSA_SUCCESS:PSA_SUCCESS
+
+PSA symmetric decrypt multipart: AES-CTR, 16 bytes, fallback
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES:MBEDTLS_PSA_BUILTIN_CIPHER
+cipher_decrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"396ee84fb75fdbb5c2b13c7fe5a654aa":"dd3b5e5319b7591daab1e1a92687feb2":0:PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+
+PSA symmetric decrypt multipart: AES-CTR, 16 bytes, fake
+depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
+cipher_decrypt_simple_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"396ee84fb75fdbb5c2b13c7fe5a654aa":"d07a6a6e2687feb2":1:PSA_SUCCESS:PSA_SUCCESS
+
 PSA symmetric encryption multipart: AES-CTR, 11+5 bytes
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_KEY_TYPE_AES
 cipher_encrypt_multipart:PSA_ALG_CTR:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":"2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a2a":"6bc1bee22e409f96e93d7e117393172a":11:11:5:"8f9408fe80a81d3e813da3c7b0b2bd32"
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 3a9eff9..7f2ac91 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -422,12 +422,159 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void cipher_encrypt( int alg_arg, int key_type_arg,
-                     data_t *key_data, data_t *iv,
-                     data_t *input, data_t *expected_output,
-                     int mock_output_arg,
-                     int force_status_arg,
-                     int expected_status_arg )
+void cipher_encrypt_sanity( int alg_arg,
+                            int key_type_arg,
+                            data_t *key_data,
+                            data_t *input,
+                            data_t *expected_output,
+                            int mock_output_arg,
+                            int force_status_arg,
+                            int expected_status_arg )
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_status_t status;
+    psa_key_type_t key_type = key_type_arg;
+    psa_algorithm_t alg = alg_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_status_t force_status = force_status_arg;
+    unsigned char *output = NULL;
+    size_t output_buffer_size = 0;
+    size_t output_length = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+    mbedtls_test_driver_cipher_hooks.forced_status = force_status;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    output_buffer_size = PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len );
+    ASSERT_ALLOC( output, output_buffer_size );
+
+    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
+                                &key ) );
+
+    if( mock_output_arg )
+    {
+        mbedtls_test_driver_cipher_hooks.forced_output = expected_output->x;
+        mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
+    }
+
+    status = psa_cipher_encrypt( key, alg, input->x, input->len, output,
+                                 output_buffer_size, &output_length );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    TEST_EQUAL( status, expected_status );
+    if( expected_status == PSA_SUCCESS && mock_output_arg )
+    {
+        ASSERT_COMPARE( expected_output->x, expected_output->len,
+                        output, output_length );
+    }
+
+exit:
+    mbedtls_free( output );
+    psa_destroy_key( key );
+    PSA_DONE( );
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void cipher_encrypt_validation( int alg_arg,
+                                int key_type_arg,
+                                data_t *key_data,
+                                data_t *input )
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_type_t key_type = key_type_arg;
+    psa_algorithm_t alg = alg_arg;
+    size_t iv_size = PSA_CIPHER_IV_LENGTH ( key_type, alg );
+    unsigned char *output1 = NULL;
+    size_t output1_buffer_size = 0;
+    size_t output1_length = 0;
+    unsigned char *output2 = NULL;
+    size_t output2_buffer_size = 0;
+    size_t output2_length = 0;
+    size_t function_output_length = 0;
+    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    output1_buffer_size = PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len );
+    output2_buffer_size = PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, input->len ) +
+                          PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg );
+    ASSERT_ALLOC( output1, output1_buffer_size );
+    ASSERT_ALLOC( output2, output2_buffer_size );
+
+    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
+                                &key ) );
+
+    PSA_ASSERT( psa_cipher_encrypt( key, alg, input->x, input->len, output1,
+                                    output1_buffer_size, &output1_length ) );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    PSA_ASSERT( psa_cipher_set_iv( &operation, output1, iv_size ) );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    PSA_ASSERT( psa_cipher_update( &operation,
+                                   input->x, input->len,
+                                   output2, output2_buffer_size,
+                                   &function_output_length ) );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    output2_length += function_output_length;
+    PSA_ASSERT( psa_cipher_finish( &operation,
+                                   output2 + output2_length,
+                                   output2_buffer_size - output2_length,
+                                   &function_output_length ) );
+    /* Finish will have called abort as well, so expecting two hits here */
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    output2_length += function_output_length;
+
+    PSA_ASSERT( psa_cipher_abort( &operation ) );
+    // driver function should've been called as part of the finish() core routine
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
+    ASSERT_COMPARE( output1 + iv_size, output1_length - iv_size,
+                    output2, output2_length );
+
+exit:
+    psa_cipher_abort( &operation );
+    mbedtls_free( output1 );
+    mbedtls_free( output2 );
+    psa_destroy_key( key );
+    PSA_DONE( );
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void cipher_encrypt_simple_multipart( int alg_arg,
+                                      int key_type_arg,
+                                      data_t *key_data,
+                                      data_t *iv,
+                                      data_t *input,
+                                      data_t *expected_output,
+                                      int mock_output_arg,
+                                      int force_status_arg,
+                                      int expected_status_arg )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
@@ -515,11 +662,14 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void cipher_encrypt_multipart( int alg_arg, int key_type_arg,
-                               data_t *key_data, data_t *iv,
+void cipher_encrypt_multipart( int alg_arg,
+                               int key_type_arg,
+                               data_t *key_data,
+                               data_t *iv,
                                data_t *input,
                                int first_part_size_arg,
-                               int output1_length_arg, int output2_length_arg,
+                               int output1_length_arg,
+                               int output2_length_arg,
                                data_t *expected_output )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -600,11 +750,14 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void cipher_decrypt_multipart( int alg_arg, int key_type_arg,
-                               data_t *key_data, data_t *iv,
+void cipher_decrypt_multipart( int alg_arg,
+                               int key_type_arg,
+                               data_t *key_data,
+                               data_t *iv,
                                data_t *input,
                                int first_part_size_arg,
-                               int output1_length_arg, int output2_length_arg,
+                               int output1_length_arg,
+                               int output2_length_arg,
                                data_t *expected_output )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -687,12 +840,15 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void cipher_decrypt( int alg_arg, int key_type_arg,
-                     data_t *key_data, data_t *iv,
-                     data_t *input, data_t *expected_output,
-                     int mock_output_arg,
-                     int force_status_arg,
-                     int expected_status_arg )
+void cipher_decrypt_simple_multipart( int alg_arg,
+                                      int key_type_arg,
+                                      data_t *key_data,
+                                      data_t *iv,
+                                      data_t *input,
+                                      data_t *expected_output,
+                                      int mock_output_arg,
+                                      int force_status_arg,
+                                      int expected_status_arg )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
@@ -779,6 +935,81 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void cipher_decrypt( int alg_arg,
+                     int key_type_arg,
+                     data_t *key_data,
+                     data_t *iv,
+                     data_t *input_arg,
+                     data_t *expected_output,
+                     int mock_output_arg,
+                     int force_status_arg,
+                     int expected_status_arg )
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_status_t status;
+    psa_key_type_t key_type = key_type_arg;
+    psa_algorithm_t alg = alg_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_status_t force_status = force_status_arg;
+    unsigned char *input = NULL;
+    size_t input_buffer_size = 0;
+    unsigned char *output = NULL;
+    size_t output_buffer_size = 0;
+    size_t output_length = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+    mbedtls_test_driver_cipher_hooks.forced_status = force_status;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    /* Allocate input buffer and copy the iv and the plaintext */
+    input_buffer_size = ( (size_t) input_arg->len + (size_t) iv->len );
+    if ( input_buffer_size > 0 )
+    {
+        ASSERT_ALLOC( input, input_buffer_size );
+        memcpy( input, iv->x, iv->len );
+        memcpy( input + iv->len, input_arg->x, input_arg->len );
+    }
+
+    output_buffer_size = PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, input_buffer_size );
+    ASSERT_ALLOC( output, output_buffer_size );
+
+    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
+                                &key ) );
+
+    if( mock_output_arg )
+    {
+        mbedtls_test_driver_cipher_hooks.forced_output = expected_output->x;
+        mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
+    }
+
+    status = psa_cipher_decrypt( key, alg, input, input_buffer_size, output,
+                                 output_buffer_size, &output_length );
+    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+
+    TEST_EQUAL( status, expected_status );
+
+    if( expected_status == PSA_SUCCESS )
+    {
+        ASSERT_COMPARE( expected_output->x, expected_output->len,
+                        output, output_length );
+    }
+
+exit:
+    mbedtls_free( input );
+    mbedtls_free( output );
+    psa_destroy_key( key );
+    PSA_DONE( );
+    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void cipher_entry_points( int alg_arg, int key_type_arg,
                           data_t *key_data, data_t *iv,
                           data_t *input )