diff --git a/TESTS/psa/entropy_inject/main.cpp b/TESTS/psa/entropy_inject/main.cpp
index 69eefa5107..314b3e5d92 100644
--- a/TESTS/psa/entropy_inject/main.cpp
+++ b/TESTS/psa/entropy_inject/main.cpp
@@ -24,7 +24,7 @@
 #include "unity/unity.h"
 #include "utest/utest.h"
 #include "psa/internal_trusted_storage.h"
-#include "test_pits.h"
+#include "psa/lifecycle.h"
 #include "entropy.h"
 #include "entropy_poll.h"
 #include "crypto.h"
@@ -134,7 +134,7 @@ static void injection_and_init_deinit()
 utest::v1::status_t case_teardown_handler(const Case *const source, const size_t passed, const size_t failed, const failure_t reason)
 {
     psa_status_t status;
-    status = test_psa_its_reset();
+    status = mbed_psa_reboot_and_request_new_security_state(PSA_LIFECYCLE_ASSEMBLY_AND_TEST);
     TEST_ASSERT_EQUAL(PSA_ITS_SUCCESS, status);
     mbedtls_psa_crypto_free();
     return greentea_case_teardown_handler(source, passed, failed, reason);
@@ -143,7 +143,7 @@ utest::v1::status_t case_teardown_handler(const Case *const source, const size_t
 utest::v1::status_t case_setup_handler(const Case *const source, const size_t index_of_case)
 {
     psa_status_t status;
-    status = test_psa_its_reset();
+    status = mbed_psa_reboot_and_request_new_security_state(PSA_LIFECYCLE_ASSEMBLY_AND_TEST);
     TEST_ASSERT_EQUAL(PSA_ITS_SUCCESS, status);
     return greentea_case_setup_handler(source, index_of_case);
 }
diff --git a/TESTS/psa/prot_internal_storage/main.cpp b/TESTS/psa/prot_internal_storage/main.cpp
index 20f729adc4..ceb2645db3 100644
--- a/TESTS/psa/prot_internal_storage/main.cpp
+++ b/TESTS/psa/prot_internal_storage/main.cpp
@@ -24,7 +24,7 @@
 #include "unity/unity.h"
 #include "utest/utest.h"
 #include "psa/internal_trusted_storage.h"
-#include "test_pits.h"
+#include "psa/lifecycle.h"
 
 using namespace utest::v1;
 
@@ -38,7 +38,7 @@ static void pits_test()
     struct psa_its_info_t info = {0, PSA_ITS_FLAG_WRITE_ONCE};
     memset(read_buff, 0, TEST_BUFF_SIZE);
 
-    status = test_psa_its_reset();
+    status = mbed_psa_reboot_and_request_new_security_state(PSA_LIFECYCLE_ASSEMBLY_AND_TEST);
     TEST_ASSERT_EQUAL(PSA_ITS_SUCCESS, status);
 
     status = psa_its_get_info(5, &info);
@@ -76,9 +76,9 @@ static void pits_write_once_test()
     psa_its_status_t status = PSA_ITS_SUCCESS;
     uint8_t write_buff[TEST_BUFF_SIZE] = {0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00};
     uint8_t read_buff[TEST_BUFF_SIZE] = {0};
-    struct psa_its_info_t info = {0, PSA_ITS_FLAG_NONE};
+    struct psa_its_info_t info = {0, 0};
 
-    status = test_psa_its_reset();
+    status = mbed_psa_reboot_and_request_new_security_state(PSA_LIFECYCLE_ASSEMBLY_AND_TEST);
     TEST_ASSERT_EQUAL(PSA_ITS_SUCCESS, status);
 
     status = psa_its_get_info(5, &info);
@@ -114,7 +114,7 @@ static void pits_write_once_test()
     TEST_ASSERT_EQUAL(TEST_BUFF_SIZE, info.size);
     TEST_ASSERT_EQUAL(PSA_ITS_FLAG_WRITE_ONCE, info.flags);
 
-    status = test_psa_its_reset();
+    status = mbed_psa_reboot_and_request_new_security_state(PSA_LIFECYCLE_ASSEMBLY_AND_TEST);
     TEST_ASSERT_EQUAL(PSA_ITS_SUCCESS, status);
 }