diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index f82eeb8..f09cc94 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -522,10 +522,7 @@
 {
     mbedtls_pk_context pk;
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-    PSA_INIT();
-#endif
-
+    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
@@ -539,9 +536,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-    PSA_DONE();
-#endif
+    USE_PSA_DONE();
 }
 /* END_CASE */