diff --git a/tests/suites/test_suite_ecjpake.function b/tests/suites/test_suite_ecjpake.function
index d5c0f07..20f1bf8 100644
--- a/tests/suites/test_suite_ecjpake.function
+++ b/tests/suites/test_suite_ecjpake.function
@@ -109,6 +109,8 @@
     mbedtls_md_type_t valid_md = MBEDTLS_MD_SHA256;
     mbedtls_ecp_group_id valid_group = MBEDTLS_ECP_DP_SECP256R1;
 
+    MD_PSA_INIT();
+
     mbedtls_ecjpake_init(&ctx);
 
     TEST_EQUAL(MBEDTLS_ERR_ECP_BAD_INPUT_DATA,
@@ -118,14 +120,19 @@
                                      valid_group,
                                      buf, len));
 exit:
-    return;
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST */
 void ecjpake_selftest()
 {
+    MD_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ecjpake_self_test(1) == 0);
+
+exit:
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -137,6 +144,8 @@
     const size_t pw_len = 0;
     int any_role = MBEDTLS_ECJPAKE_CLIENT;
 
+    MD_PSA_INIT();
+
     mbedtls_ecjpake_init(&corrupt_ctx);
     TEST_ASSERT(mbedtls_ecjpake_setup(&corrupt_ctx, any_role,
                                       MBEDTLS_MD_SHA256, MBEDTLS_ECP_DP_SECP256R1, pw,
@@ -148,6 +157,7 @@
 
 exit:
     mbedtls_ecjpake_free(&corrupt_ctx);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -158,6 +168,8 @@
     const unsigned char *pw = NULL;
     const size_t pw_len = 0;
 
+    MD_PSA_INIT();
+
     mbedtls_ecjpake_init(&ctx);
 
     TEST_ASSERT(mbedtls_ecjpake_setup(&ctx, role,
@@ -168,6 +180,7 @@
 
 exit:
     mbedtls_ecjpake_free(&ctx);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -178,6 +191,8 @@
     const unsigned char *pw = NULL;
     const size_t pw_len = 0;
 
+    MD_PSA_INIT();
+
     mbedtls_ecjpake_init(&ctx);
 
     TEST_ASSERT(mbedtls_ecjpake_setup(&ctx, MBEDTLS_ECJPAKE_CLIENT,
@@ -194,6 +209,7 @@
 
 exit:
     mbedtls_ecjpake_free(&ctx);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -204,6 +220,8 @@
     const unsigned char *pw = NULL;
     const size_t pw_len = 0;
 
+    MD_PSA_INIT();
+
     mbedtls_ecjpake_init(&ctx);
 
     TEST_ASSERT(mbedtls_ecjpake_setup(&ctx, MBEDTLS_ECJPAKE_SERVER,
@@ -220,5 +238,6 @@
 
 exit:
     mbedtls_ecjpake_free(&ctx);
+    MD_PSA_DONE();
 }
 /* END_CASE */