Move the setting of id in driver wrappers

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 354477a..48410c0 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -1458,26 +1458,25 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
+            operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
             status = mbedtls_test_transparent_aead_encrypt_setup(
                         &operation->ctx.transparent_test_driver_ctx,
                         attributes, key_buffer, key_buffer_size,
                         alg );
-            /* Declared with fallback == true */
-            operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
 
+            /* Declared with fallback == true */
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 return( status );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
 
             /* Fell through, meaning no accelerator supports this operation */
+            operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
             status = mbedtls_psa_aead_encrypt_setup(
                         &operation->ctx.mbedtls_ctx, attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
-            operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
-
             return( status );
 
         /* Add cases for opaque driver here */
@@ -1507,28 +1506,27 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
+            operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
             status = mbedtls_test_transparent_aead_decrypt_setup(
                         &operation->ctx.transparent_test_driver_ctx,
                          attributes,
                         key_buffer, key_buffer_size,
                         alg );
-            /* Declared with fallback == true */
-            operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
 
+            /* Declared with fallback == true */
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 return( status );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
 
             /* Fell through, meaning no accelerator supports this operation */
+            operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
             status = mbedtls_psa_aead_decrypt_setup(
                         &operation->ctx.mbedtls_ctx,
                          attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
-            operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
-
             return( status );
 
         /* Add cases for opaque driver here */