tests: add PSA_INIT/PSA_DONE to CCM and GCM test suites

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 0b8c221..41d204d 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -367,6 +367,30 @@
 #define MD_PSA_DONE() ((void) 0)
 #endif /* MBEDTLS_MD_SOME_PSA */
 
+/** \def BLOCK_CIPHER_PSA_INIT
+ *
+ * Call this macro to initialize the PSA subsystem if BLOCK_CIPHER uses a driver,
+ * and do nothing otherwise.
+ *
+ * If the initialization fails, mark the test case as failed and jump to the
+ * \p exit label.
+ */
+/** \def BLOCK_CIPHER_PSA_DONE
+ *
+ * Call this macro at the end of a test case if you called #BLOCK_CIPHER_PSA_INIT.
+ *
+ * This is like #PSA_DONE except it does nothing under the same conditions as
+ * #BLOCK_CIPHER_PSA_INIT.
+ */
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+#define BLOCK_CIPHER_PSA_INIT()   PSA_INIT()
+#define BLOCK_CIPHER_PSA_DONE()   PSA_DONE()
+#else /* MBEDTLS_MD_SOME_PSA */
+#define BLOCK_CIPHER_PSA_INIT() ((void) 0)
+#define BLOCK_CIPHER_PSA_DONE() ((void) 0)
+#endif /* MBEDTLS_MD_SOME_PSA */
+
+
 /** \def MD_OR_USE_PSA_INIT
  *
  * Call this macro to initialize the PSA subsystem if MD uses a driver,
diff --git a/tests/suites/test_suite_ccm.function b/tests/suites/test_suite_ccm.function
index 5aaaaa2..0685e5b 100644
--- a/tests/suites/test_suite_ccm.function
+++ b/tests/suites/test_suite_ccm.function
@@ -69,7 +69,9 @@
 /* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST:MBEDTLS_AES_C */
 void mbedtls_ccm_self_test()
 {
+    BLOCK_CIPHER_PSA_INIT();
     TEST_ASSERT(mbedtls_ccm_self_test(1) == 0);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -105,6 +107,7 @@
     unsigned char tag[18];
     int decrypt_ret;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
 
     TEST_CALLOC_OR_SKIP(add, add_len);
@@ -132,6 +135,7 @@
 exit:
     mbedtls_free(add);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -148,6 +152,7 @@
     unsigned char tag[18];
     int decrypt_ret;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
 
     memset(key, 0, sizeof(key));
@@ -174,6 +179,7 @@
 
 exit:
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -198,6 +204,7 @@
     /* Prepare tag buffer */
     TEST_CALLOC(tag_buf, expected_tag_len);
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     /* Test with input == output */
@@ -230,6 +237,7 @@
     mbedtls_ccm_free(&ctx);
     mbedtls_free(io_msg_buf);
     mbedtls_free(tag_buf);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -241,6 +249,7 @@
     uint8_t *output = NULL;
     size_t olen;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -255,6 +264,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -277,6 +287,7 @@
         memcpy(io_msg_buf, msg->x, expected_msg_len);
     }
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     /* Test with input == output */
@@ -317,6 +328,7 @@
 exit:
     mbedtls_free(io_msg_buf);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -364,6 +376,7 @@
     iv[source_address->len + frame_counter->len] = sec_level;
     iv_len = sizeof(iv);
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id,
                                   key->x, key->len * 8), 0);
@@ -402,6 +415,7 @@
     mbedtls_ccm_free(&ctx);
     mbedtls_free(io_msg_buf);
     mbedtls_free(tag_buf);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -443,6 +457,7 @@
     iv[source_address->len + frame_counter->len] = sec_level;
     iv_len = sizeof(iv);
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_ASSERT(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8) == 0);
     /* Test with input == output */
@@ -479,6 +494,7 @@
 exit:
     mbedtls_ccm_free(&ctx);
     mbedtls_free(io_msg_buf);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -495,6 +511,7 @@
     /* Sanity checks on the test data */
     TEST_EQUAL(msg->len, result->len);
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -517,6 +534,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -529,6 +547,7 @@
     mbedtls_ccm_context ctx;
     uint8_t *output = NULL;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -545,6 +564,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -556,6 +576,7 @@
 {
     mbedtls_ccm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -566,6 +587,7 @@
     TEST_EQUAL(MBEDTLS_ERR_CCM_BAD_INPUT, mbedtls_ccm_update_ad(&ctx, add->x, add->len));
 exit:
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -577,6 +599,7 @@
 {
     mbedtls_ccm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -586,6 +609,7 @@
     TEST_EQUAL(MBEDTLS_ERR_CCM_BAD_INPUT, mbedtls_ccm_update_ad(&ctx, add->x, add->len));
 exit:
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -599,6 +623,7 @@
     uint8_t *output = NULL;
     size_t olen;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -614,6 +639,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -625,6 +651,7 @@
     mbedtls_ccm_context ctx;
     uint8_t *output = NULL;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -639,6 +666,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -651,6 +679,7 @@
 {
     mbedtls_ccm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -663,6 +692,7 @@
     TEST_EQUAL(MBEDTLS_ERR_CCM_BAD_INPUT, mbedtls_ccm_update_ad(&ctx, add->x, 1));
 exit:
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -679,6 +709,7 @@
     add_second_buffer[0] = add->x[add->len - 1];
     add_second_buffer[1] = 0xAB; // some magic value
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -691,6 +722,7 @@
     TEST_EQUAL(MBEDTLS_ERR_CCM_BAD_INPUT, mbedtls_ccm_update_ad(&ctx, add_second_buffer, 2));
 exit:
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -704,6 +736,7 @@
     uint8_t *output = NULL;
     size_t olen;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -719,6 +752,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -732,6 +766,7 @@
     uint8_t *output = NULL;
     size_t olen;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -752,6 +787,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -766,6 +802,7 @@
     uint8_t *output = NULL;
     size_t olen;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -783,6 +820,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -801,6 +839,7 @@
     msg_second_buffer[0] = msg->x[msg->len - 1];
     msg_second_buffer[1] = 0xAB; // some magic value
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -818,6 +857,7 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -829,6 +869,7 @@
     mbedtls_ccm_context ctx;
     uint8_t *output = NULL;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_ccm_init(&ctx);
     TEST_EQUAL(mbedtls_ccm_setkey(&ctx, cipher_id, key->x, key->len * 8), 0);
     TEST_EQUAL(0, mbedtls_ccm_starts(&ctx, mode, iv->x, iv->len));
@@ -842,5 +883,6 @@
 exit:
     mbedtls_free(output);
     mbedtls_ccm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_gcm.function b/tests/suites/test_suite_gcm.function
index 747914f..097e424 100644
--- a/tests/suites/test_suite_gcm.function
+++ b/tests/suites/test_suite_gcm.function
@@ -171,6 +171,7 @@
     mbedtls_gcm_context ctx;
     size_t tag_len = tag_len_bits / 8;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     memset(output, 0x00, sizeof(output));
@@ -183,6 +184,7 @@
 
 exit:
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -200,6 +202,7 @@
     size_t n1;
     size_t n1_add;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     memset(output, 0x00, 128);
@@ -230,6 +233,7 @@
 
 exit:
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -247,6 +251,7 @@
     size_t n1;
     size_t n1_add;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     memset(output, 0x00, 128);
@@ -287,6 +292,7 @@
 
 exit:
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -300,6 +306,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -308,6 +315,7 @@
                                cipher_update_calls);
 
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -322,6 +330,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -330,6 +339,7 @@
                                ad_update_calls);
 
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -341,6 +351,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -348,6 +359,7 @@
                           iv_str, tag_str);
 
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -361,6 +373,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -370,6 +383,7 @@
 
 exit:
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -384,6 +398,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -393,6 +408,7 @@
 
 exit:
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -404,6 +420,7 @@
 {
     mbedtls_gcm_context ctx;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
 
     TEST_ASSERT(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8) == 0);
@@ -411,6 +428,7 @@
                           iv_str, tag_str);
 
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
@@ -444,6 +462,7 @@
     size_t olen = 0;
     size_t output_len = input->len - 1;
 
+    BLOCK_CIPHER_PSA_INIT();
     mbedtls_gcm_init(&ctx);
     TEST_EQUAL(mbedtls_gcm_setkey(&ctx, cipher_id, key_str->x, key_str->len * 8), 0);
     TEST_EQUAL(0, mbedtls_gcm_starts(&ctx, mode, iv->x, iv->len));
@@ -455,12 +474,15 @@
 exit:
     mbedtls_free(output);
     mbedtls_gcm_free(&ctx);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST:MBEDTLS_AES_C */
 void gcm_selftest()
 {
+    BLOCK_CIPHER_PSA_INIT();
     TEST_ASSERT(mbedtls_gcm_self_test(1) == 0);
+    BLOCK_CIPHER_PSA_DONE();
 }
 /* END_CASE */