test: proper positioning of USE_PSA_INIT + added missing exit labels

Signed-off-by: valerio <valerio.setti@nordicsemi.no>
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);