diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index faa7dbe..684a90a 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -216,10 +216,9 @@
     size_t len;
     mbedtls_pk_debug_item dbg;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_pk_init(&pk2);
+    USE_PSA_INIT();
 
     TEST_ASSERT(psa_crypto_init() == PSA_SUCCESS);
 
@@ -314,9 +313,8 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     if (opaque_key == 1) {
         psa_set_key_usage_flags(&attributes, key_usage);
@@ -361,9 +359,8 @@
     unsigned char buf[] = { 0x01, 0x02, 0x03, 0x04, 0x05, 0x06 };
     size_t buf_size = sizeof(buf);
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
                mbedtls_pk_verify_restartable(&ctx, MBEDTLS_MD_NONE,
@@ -411,8 +408,8 @@
     size_t len;
     void *options = NULL;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
@@ -497,9 +494,10 @@
 {
     mbedtls_pk_context pk;
 
-    USE_PSA_INIT();
     /* For the write tests to be effective, we need a valid key pair. */
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_parse_key(&pk,
                                      key_data->x, key_data->len, NULL, 0,
                                      mbedtls_test_rnd_std_rand, NULL) == 0);
@@ -529,8 +527,8 @@
 {
     mbedtls_pk_context pk;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
     TEST_ASSERT(pk_genkey(&pk, parameter) == 0);
@@ -552,11 +550,10 @@
 {
     mbedtls_pk_context pub, prv, alt;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pub);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&alt);
+    USE_PSA_INIT();
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     /* mbedtls_pk_check_pair() returns either PK or ECP error codes depending
@@ -611,10 +608,8 @@
     mbedtls_ecp_set_max_ops(1);
 #endif
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
-
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -653,8 +648,8 @@
     void *options;
     int ret;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -708,8 +703,8 @@
     mbedtls_pk_context pk;
     mbedtls_ecp_keypair *eckey;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
 
@@ -742,11 +737,11 @@
     unsigned char sig[MBEDTLS_ECDSA_MAX_LEN];
     size_t slen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_restart_init(&rs_ctx);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&pub);
+    USE_PSA_INIT();
+
     memset(sig, 0, sizeof(sig));
 
     TEST_ASSERT(mbedtls_pk_setup(&prv, mbedtls_pk_info_from_type(pk_type)) == 0);
@@ -837,8 +832,8 @@
     mbedtls_ecp_set_max_ops(42000);
 #endif
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    MD_OR_USE_PSA_INIT();
 
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
@@ -919,11 +914,10 @@
     mbedtls_pk_context pk;
     size_t olen, rlen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
     memset(output,     0, sizeof(output));
@@ -995,15 +989,13 @@
     mbedtls_pk_context pk;
     size_t olen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
 
-
     /* init pk-rsa context */
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -1052,11 +1044,10 @@
     mbedtls_svc_key_id_t key_id;
     size_t olen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
 
@@ -1113,8 +1104,8 @@
     size_t olen = 0;
     int ret = MBEDTLS_ERR_PK_TYPE_MISMATCH;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
     memset(output,     0, sizeof(output));
@@ -1143,12 +1134,12 @@
     size_t hash_len = SIZE_MAX, sig_len = SIZE_MAX;
     unsigned char hash[50], sig[100];
 
+    mbedtls_pk_init(&pk);
     USE_PSA_INIT();
+
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    mbedtls_pk_init(&pk);
-
     TEST_ASSERT(mbedtls_pk_setup(&pk,
                                  mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
 
@@ -1188,10 +1179,10 @@
     size_t sig_len, ciph_len, test_len;
     int ret = MBEDTLS_ERR_PK_TYPE_MISMATCH;
 
-    USE_PSA_INIT();
-
     mbedtls_rsa_init(&raw);
-    mbedtls_pk_init(&rsa); mbedtls_pk_init(&alt);
+    mbedtls_pk_init(&rsa);
+    mbedtls_pk_init(&alt);
+    USE_PSA_INIT();
 
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
@@ -1285,12 +1276,12 @@
      * - parse it to a PK context and verify the signature this way
      */
 
+    mbedtls_pk_init(&pk);
     USE_PSA_INIT();
 
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if (PSA_KEY_TYPE_IS_RSA(psa_type_arg)) {
         /* Create legacy RSA public/private key in PK context. */
-        mbedtls_pk_init(&pk);
         TEST_ASSERT(mbedtls_pk_setup(&pk,
                                      mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
         TEST_ASSERT(mbedtls_rsa_gen_key(mbedtls_pk_rsa(pk),
@@ -1304,7 +1295,6 @@
         mbedtls_ecp_group_id grpid = parameter_arg;
 
         /* Create legacy EC public/private key in PK context. */
-        mbedtls_pk_init(&pk);
         TEST_ASSERT(mbedtls_pk_setup(&pk,
                                      mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
         TEST_ASSERT(pk_genkey(&pk, grpid) == 0);
@@ -1424,8 +1414,8 @@
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    PSA_INIT();
     mbedtls_pk_init(&pk);
+    PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk,
                                  mbedtls_pk_info_from_type(pk_type)) == 0);
@@ -1468,8 +1458,8 @@
     mbedtls_pk_rsassa_pss_options rsassa_pss_options;
     int ret;
 
-    PSA_INIT();
     mbedtls_pk_init(&pk);
+    PSA_INIT();
 
     /* Create legacy RSA public/private key in PK context. */
     mbedtls_pk_init(&pk);
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 5d8a7c8..7e05304 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -1118,10 +1118,9 @@
     mbedtls_ssl_context ssl;
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_config_defaults(&conf,
                                             MBEDTLS_SSL_IS_CLIENT,
@@ -1151,12 +1150,13 @@
 {
     mbedtls_ssl_context ssl;
 
-    USE_PSA_INIT();
     mbedtls_ssl_init(&ssl);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_set_hostname(&ssl, hostname0) == 0);
     TEST_ASSERT(mbedtls_ssl_set_hostname(&ssl, hostname1) == 0);
 
+exit:
     mbedtls_ssl_free(&ssl);
     USE_PSA_DONE();
 }
@@ -1182,11 +1182,11 @@
     size_t const buflen = 512;
     mbedtls_record rec, rec_backup;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
+
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
                                             etm, tag_mode, ver,
                                             (size_t) cid0_len,
@@ -1336,11 +1336,11 @@
     int seen_success; /* Indicates if in the current mode we've
                        * already seen a successful test. */
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
+
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
                                             etm, tag_mode, ver,
                                             (size_t) cid0_len,
@@ -1497,11 +1497,10 @@
     int ret;
     const unsigned char pad_max_len = 255; /* Per the standard */
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
 
     /* Set up transforms with dummy keys */
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
@@ -1726,6 +1725,7 @@
     ASSERT_COMPARE(dst, (size_t) desired_length,
                    expected->x, (size_t) expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1779,6 +1779,7 @@
                    expected_server_write_iv->x,
                    (size_t) desired_iv_len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1823,6 +1824,7 @@
     ASSERT_COMPARE(dst, desired_length,
                    expected->x, desired_length);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1856,6 +1858,7 @@
     ASSERT_COMPARE(secrets.early_exporter_master_secret, hash_len,
                    exporter_expected->x, exporter_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1889,6 +1892,7 @@
     ASSERT_COMPARE(secrets.server_handshake_traffic_secret, hash_len,
                    server_expected->x, server_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1926,6 +1930,7 @@
     ASSERT_COMPARE(secrets.exporter_master_secret, hash_len,
                    exporter_expected->x, exporter_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1955,6 +1960,7 @@
     ASSERT_COMPARE(secrets.resumption_master_secret, hash_len,
                    resumption_expected->x, resumption_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1988,6 +1994,7 @@
     ASSERT_COMPARE(binder, hash_len,
                    binder_expected->x, binder_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -2086,6 +2093,7 @@
     ASSERT_COMPARE(rec.buf + rec.data_offset, rec.data_len,
                    plaintext->x, plaintext->len);
 
+exit:
     mbedtls_free(buf);
     mbedtls_ssl_transform_free(&transform_send);
     mbedtls_ssl_transform_free(&transform_recv);
@@ -2112,6 +2120,7 @@
     ASSERT_COMPARE(secret_new, (size_t) expected->len,
                    expected->x, (size_t) expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -2155,9 +2164,9 @@
     /*
      * Test that a save-load pair is the identity
      */
-    USE_PSA_INIT();
     mbedtls_ssl_session_init(&original);
     mbedtls_ssl_session_init(&restored);
+    USE_PSA_INIT();
 
     /* Prepare a dummy session to work on */
     ((void) endpoint_type);
@@ -2304,8 +2313,8 @@
     /*
      * Test that a load-save pair is the identity
      */
-    USE_PSA_INIT();
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare a dummy session to work on */
     ((void) endpoint_type);
@@ -2367,8 +2376,8 @@
     /*
      * Test that session_save() fails cleanly on small buffers
      */
-    USE_PSA_INIT();
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare dummy session and get serialized size */
     ((void) endpoint_type);
@@ -2415,8 +2424,8 @@
     /*
      * Test that session_load() fails cleanly on small buffers
      */
-    USE_PSA_INIT();
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare serialized session data */
     ((void) endpoint_type);
@@ -2476,8 +2485,8 @@
                                       corrupt_config == 1,
                                       corrupt_config == 1 };
 
-    USE_PSA_INIT();
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
     ((void) endpoint_type);
     ((void) tls_version);
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
@@ -2530,6 +2539,7 @@
             *byte ^= corrupted_bit;
         }
     }
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
@@ -2697,6 +2707,8 @@
 
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+
+exit:
     MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
@@ -2906,9 +2918,11 @@
     USE_PSA_INIT();
     test_resize_buffers(mfl, 0, MBEDTLS_SSL_LEGACY_NO_RENEGOTIATION, 1, 1,
                         (char *) "");
-    USE_PSA_DONE();
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2918,9 +2932,11 @@
 {
     USE_PSA_INIT();
     test_resize_buffers(mfl, 1, legacy_renegotiation, 0, 1, cipher);
-    USE_PSA_DONE();
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2935,8 +2951,8 @@
 
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_conf_psk(&conf,
                                      psk0, sizeof(psk0),
@@ -2947,9 +2963,7 @@
                 MBEDTLS_ERR_SSL_FEATURE_UNAVAILABLE);
 
 exit:
-
     mbedtls_ssl_config_free(&conf);
-
     MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
@@ -2979,8 +2993,8 @@
 
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     switch (mode) {
         case 0:
@@ -3031,7 +3045,6 @@
     }
 
 exit:
-
     mbedtls_ssl_config_free(&conf);
     MD_OR_USE_PSA_DONE();
 
@@ -3046,10 +3059,9 @@
     mbedtls_ssl_config conf;
     mbedtls_ssl_context ssl;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_config_init(&conf);
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_conf_endpoint(&conf, endpoint);
     mbedtls_ssl_conf_transport(&conf, transport);
@@ -3090,10 +3102,10 @@
 #endif
     mbedtls_ssl_conf_curves(&conf, curve_list);
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
 
     TEST_ASSERT(ssl.handshake != NULL && ssl.handshake->group_list != NULL);
@@ -3107,7 +3119,6 @@
         TEST_EQUAL(iana_tls_group_list[i], ssl.handshake->group_list[i]);
     }
 
-
 exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
@@ -3131,10 +3142,10 @@
 
     mbedtls_ssl_conf_groups(&conf, iana_tls_group_list);
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
 
     TEST_ASSERT(ssl.conf != NULL && ssl.conf->group_list != NULL);
@@ -3227,9 +3238,10 @@
     mbedtls_ssl_config conf;
     size_t len;
 
-    USE_PSA_INIT();
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    USE_PSA_INIT();
+
     TEST_EQUAL(mbedtls_ssl_config_defaults(&conf, MBEDTLS_SSL_IS_SERVER,
                                            MBEDTLS_SSL_TRANSPORT_DATAGRAM,
                                            MBEDTLS_SSL_PRESET_DEFAULT),
@@ -3244,6 +3256,7 @@
                                                     &len),
                exp_ret);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     USE_PSA_DONE();
@@ -3259,6 +3272,8 @@
     mbedtls_timing_set_delay(&delay_context, 50, 100);
 
     TEST_ASSERT(mbedtls_timing_get_final_delay(&delay_context) == 100);
+
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
@@ -3276,10 +3291,9 @@
 
     mbedtls_test_rnd_std_rand(NULL, own_cid, sizeof(own_cid));
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_config_defaults(&conf,
                                             MBEDTLS_SSL_IS_CLIENT,
@@ -3435,11 +3449,12 @@
     /*
      * Test set-up
      */
-    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client_ep, sizeof(client_ep));
     mbedtls_platform_zeroize(&server_ep, sizeof(server_ep));
 
     mbedtls_test_init_handshake_options(&client_options);
+    MD_OR_USE_PSA_INIT();
+
     client_options.pk_alg = MBEDTLS_PK_ECDSA;
     ret = mbedtls_test_ssl_endpoint_init(&client_ep, MBEDTLS_SSL_IS_CLIENT,
                                          &client_options, NULL, NULL, NULL,
@@ -3530,9 +3545,8 @@
     size_t pwd_len = 0;
     int ret;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
 
     /* test with uninitalized SSL context */
     ECJPAKE_TEST_SET_PASSWORD(MBEDTLS_ERR_SSL_BAD_INPUT_DATA);
@@ -3669,7 +3683,8 @@
 #else
     TEST_UNAVAILABLE_ECC(30, MBEDTLS_ECP_DP_CURVE448, PSA_ECC_FAMILY_MONTGOMERY, 448);
 #endif
-
+    goto exit;
+exit:
     MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509parse.function b/tests/suites/test_suite_x509parse.function
index 5b93b12..5abd085 100644
--- a/tests/suites/test_suite_x509parse.function
+++ b/tests/suites/test_suite_x509parse.function
@@ -428,13 +428,14 @@
     mbedtls_x509_crt crt;
     int expected_result = ext_type & has_ext_type;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     crt.ext_types = ext_type;
 
     TEST_ASSERT(mbedtls_x509_crt_has_ext_type(&crt, has_ext_type) == expected_result);
 
+exit:
     mbedtls_x509_crt_free(&crt);
     USE_PSA_DONE();
 }
@@ -464,8 +465,8 @@
     char *p = buf;
     size_t n = sizeof(buf);
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_EQUAL(mbedtls_x509_crt_parse_file(&crt, crt_file), parse_result);
@@ -493,7 +494,6 @@
     TEST_ASSERT(strcmp(buf, result_str) == 0);
 
 exit:
-
     mbedtls_x509_crt_free(&crt);
     USE_PSA_DONE();
 }
@@ -506,8 +506,8 @@
     char buf[2000];
     int res;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -531,8 +531,8 @@
     char buf[2000];
     int res;
 
-    USE_PSA_INIT();
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == 0);
@@ -555,8 +555,8 @@
     mbedtls_x509_crl   crl;
     char buf[2000];
 
-    USE_PSA_INIT();
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == result);
@@ -574,8 +574,8 @@
     char buf[2000];
     int res;
 
-    USE_PSA_INIT();
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_csr_parse_file(&csr, csr_file) == 0);
@@ -606,6 +606,8 @@
     TEST_ASSERT(res >= 0);
 
     TEST_ASSERT(strcmp(buf, result_str) == 0);
+
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
@@ -630,10 +632,10 @@
      * - x509_verify() for server5 -> test-ca2:             ~ 18800
      * - x509_verify() for server10 -> int-ca3 -> int-ca2:  ~ 25500
      */
-    MD_OR_USE_PSA_INIT();
     mbedtls_x509_crt_restart_init(&rs_ctx);
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&ca, ca_file) == 0);
@@ -682,10 +684,10 @@
     char *cn_name = NULL;
     const mbedtls_x509_crt_profile *profile;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
     mbedtls_x509_crl_init(&crl);
+    MD_OR_USE_PSA_INIT();
 
     if (strcmp(cn_name_str, "NULL") != 0) {
         cn_name = cn_name_str;
@@ -767,9 +769,9 @@
     mbedtls_x509_crt ca;
     uint32_t flags = 0;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&ca, ca_file) == 0);
@@ -801,9 +803,10 @@
     uint32_t flags = 0;
     verify_print_context vrfy_ctx;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
+    MD_OR_USE_PSA_INIT();
+
     verify_print_init(&vrfy_ctx);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -838,8 +841,9 @@
     char buf[2000];
     int res = 0;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
+
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -868,8 +872,9 @@
     char buf[2000];
     int res = 0;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
+
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -914,6 +919,7 @@
 
     TEST_EQUAL(ret, exp_ret);
 
+exit:
     mbedtls_free(name);
     USE_PSA_DONE();
 }
@@ -992,8 +998,8 @@
 {
     mbedtls_x509_crt   crt;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1016,8 +1022,8 @@
 {
     mbedtls_x509_crt   crt;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1040,8 +1046,8 @@
 {
     mbedtls_x509_crt crt;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == result);
 
@@ -1062,8 +1068,8 @@
     ((void) result_str);
 #endif
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_der(&crt, buf->x, buf->len) == (result));
 #if !defined(MBEDTLS_X509_REMOVE_INFO)
@@ -1147,12 +1153,12 @@
     ((void) result_str);
 #endif
 
-    USE_PSA_INIT();
     oid.tag = MBEDTLS_ASN1_OID;
     oid.len = MBEDTLS_OID_SIZE(MBEDTLS_OID_PKIX "\x01\x1F");
     oid.p = (unsigned char *) MBEDTLS_OID_PKIX "\x01\x1F";
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_der_with_ext_cb(&crt, buf->x, buf->len, 0, parse_crt_ext_cb,
                                                        &oid) == (result));
@@ -1197,8 +1203,9 @@
     unsigned char output[2000];
     int res;
 
-    USE_PSA_INIT();
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
+
     memset(output, 0, 2000);
 
 
@@ -1225,8 +1232,9 @@
     char my_out[1000];
     int my_ret;
 
-    USE_PSA_INIT();
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
+
     memset(my_out, 0, sizeof(my_out));
 
     my_ret = mbedtls_x509_csr_parse_der(&csr, csr_der->x, csr_der->len);
@@ -1251,8 +1259,9 @@
     char my_out[1000];
     int my_ret;
 
-    USE_PSA_INIT();
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
+
     memset(my_out, 0, sizeof(my_out));
 
     my_ret = mbedtls_x509_csr_parse_file(&csr, csr_file);
@@ -1276,8 +1285,8 @@
     mbedtls_x509_crt chain, *cur;
     int i;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&chain);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_path(&chain, crt_path) == ret);
 
@@ -1309,9 +1318,9 @@
      * We expect chain_dir to contain certificates 00.crt, 01.crt, etc.
      * with NN.crt signed by NN-1.crt
      */
-    MD_OR_USE_PSA_INIT();
     mbedtls_x509_crt_init(&trusted);
     mbedtls_x509_crt_init(&chain);
+    MD_OR_USE_PSA_INIT();
 
     /* Load trusted root */
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&trusted, ca_file) == 0);
@@ -1347,9 +1356,9 @@
     mbedtls_x509_crt trusted, chain;
     const mbedtls_x509_crt_profile *profile = NULL;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_x509_crt_init(&chain);
     mbedtls_x509_crt_init(&trusted);
+    MD_OR_USE_PSA_INIT();
 
     while ((act = mystrsep(&chain_paths, " ")) != NULL) {
         TEST_ASSERT(mbedtls_x509_crt_parse_file(&chain, act) == 0);
@@ -1389,6 +1398,7 @@
     int ret;
 
     USE_PSA_INIT();
+
     oid.tag = MBEDTLS_ASN1_OID;
     oid.p   = buf->x;
     oid.len   = buf->len;
@@ -1403,6 +1413,8 @@
         TEST_ASSERT(desc != NULL);
         TEST_ASSERT(strcmp(desc, ref_desc) == 0);
     }
+
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
@@ -1414,6 +1426,7 @@
     char num_buf[100];
 
     USE_PSA_INIT();
+
     memset(num_buf, 0x2a, sizeof(num_buf));
 
     oid.tag = MBEDTLS_ASN1_OID;
@@ -1428,6 +1441,8 @@
         TEST_ASSERT(num_buf[ret] == 0);
         TEST_ASSERT(strcmp(num_buf, numstr) == 0);
     }
+
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
@@ -1437,8 +1452,8 @@
 {
     mbedtls_x509_crt crt;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1456,8 +1471,8 @@
 {
     mbedtls_x509_crt crt;
 
-    USE_PSA_INIT();
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1497,6 +1512,7 @@
         TEST_ASSERT(min  == time.min);
         TEST_ASSERT(sec  == time.sec);
     }
+exit:
     USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index 4b0286b..b08555c 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -159,8 +159,6 @@
     const char *san_dns_name = "example.com";
     const char *san_uri_name = "http://pki.example.com/";
 
-    MD_OR_USE_PSA_INIT();
-
     san_uri.node.type = MBEDTLS_X509_SAN_UNIFORM_RESOURCE_IDENTIFIER;
     san_uri.node.san.unstructured_name.p = (unsigned char *) san_uri_name;
     san_uri.node.san.unstructured_name.len = strlen(san_uri_name);
@@ -178,8 +176,9 @@
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 
     mbedtls_x509write_csr_init(&req);
-
     mbedtls_pk_init(&key);
+    MD_OR_USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&key, key_file, NULL,
                                          mbedtls_test_rnd_std_rand, NULL) == 0);
 
@@ -266,12 +265,11 @@
     const char *subject_name = "C=NL,O=PolarSSL,CN=PolarSSL Server 1";
     mbedtls_test_rnd_pseudo_info rnd_info;
 
+    mbedtls_x509write_csr_init(&req);
     MD_OR_USE_PSA_INIT();
 
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 
-    mbedtls_x509write_csr_init(&req);
-
     md_alg_psa = mbedtls_hash_info_psa_from_md((mbedtls_md_type_t) md_type);
     TEST_ASSERT(md_alg_psa != MBEDTLS_MD_NONE);
 
@@ -351,8 +349,6 @@
 #endif
     mbedtls_pk_type_t issuer_key_type;
 
-    MD_OR_USE_PSA_INIT();
-
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 #if defined(MBEDTLS_TEST_DEPRECATED) && defined(MBEDTLS_BIGNUM_C)
     mbedtls_mpi_init(&serial_mpi);
@@ -361,8 +357,8 @@
     mbedtls_pk_init(&subject_key);
     mbedtls_pk_init(&issuer_key);
     mbedtls_pk_init(&issuer_key_alt);
-
     mbedtls_x509write_crt_init(&crt);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&subject_key, subject_key_file,
                                          subject_pwd, mbedtls_test_rnd_std_rand, NULL) == 0);