test: fix USE_PSA_INIT/DONE for x509 test suite

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index cb3f6a5..b4509e2 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -145,8 +145,9 @@
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 
     mbedtls_x509write_csr_init(&req);
-
     mbedtls_pk_init(&key);
+    USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&key, key_file, NULL) == 0);
 
     mbedtls_x509write_csr_set_md_alg(&req, md_type);
@@ -197,6 +198,7 @@
 exit:
     mbedtls_x509write_csr_free(&req);
     mbedtls_pk_free(&key);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -246,12 +248,11 @@
     buf[pem_len] = '\0';
     TEST_ASSERT(x509_crt_verifycsr(buf, pem_len + 1) == 0);
 
-
 exit:
     mbedtls_x509write_csr_free(&req);
     mbedtls_pk_free(&key);
     psa_destroy_key(key_id);
-    PSA_DONE();
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -287,6 +288,7 @@
     mbedtls_pk_init(&issuer_key_alt);
 
     mbedtls_x509write_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&subject_key, subject_key_file,
                                          subject_pwd) == 0);
@@ -422,6 +424,7 @@
     mbedtls_pk_free(&subject_key);
     mbedtls_pk_free(&issuer_key);
     mbedtls_mpi_free(&serial);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -432,6 +435,7 @@
     mbedtls_mpi serial_mpi;
     uint8_t invalid_serial[MBEDTLS_X509_RFC5280_MAX_SERIAL_LEN + 1];
 
+    USE_PSA_INIT();
     memset(invalid_serial, 0x01, sizeof(invalid_serial));
 
     mbedtls_mpi_init(&serial_mpi);
@@ -442,6 +446,7 @@
 
 exit:
     mbedtls_mpi_free(&serial_mpi);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -455,6 +460,8 @@
     mbedtls_x509_name parsed, *parsed_cur, *parsed_prv;
     unsigned char buf[1024], out[1024], *c;
 
+    USE_PSA_INIT();
+
     memset(&parsed, 0, sizeof(parsed));
     memset(out, 0, sizeof(out));
     memset(buf, 0, sizeof(buf));
@@ -488,5 +495,6 @@
         parsed_cur = parsed_cur->next;
         mbedtls_free(parsed_prv);
     }
+    USE_PSA_DONE();
 }
 /* END_CASE */