diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 2e2606f..acc2f8c 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1259,7 +1259,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1296,7 +1296,7 @@
 
 exit:
     mbedtls_free( buffer );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1389,7 +1389,7 @@
     mbedtls_free( exported );
     mbedtls_free( reexported );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1400,7 +1400,7 @@
     test_operations_on_invalid_handle( handle );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1453,7 +1453,7 @@
     mbedtls_free( exported );
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1495,7 +1495,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1527,7 +1527,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1607,7 +1607,7 @@
 exit:
     psa_mac_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1650,7 +1650,7 @@
 exit:
     psa_cipher_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1711,7 +1711,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1772,7 +1772,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( buffer );
 }
 /* END_CASE */
@@ -1827,7 +1827,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1866,7 +1866,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1904,7 +1904,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1941,7 +1941,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1978,7 +1978,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2059,7 +2059,7 @@
 exit:
     psa_reset_key_attributes( &source_attributes );
     psa_reset_key_attributes( &target_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 }
 /* END_CASE */
@@ -2103,7 +2103,7 @@
 exit:
     psa_reset_key_attributes( &source_attributes );
     psa_reset_key_attributes( &target_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2165,7 +2165,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2253,7 +2253,7 @@
     PSA_ASSERT( psa_hash_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2288,7 +2288,7 @@
                 PSA_ERROR_INVALID_SIGNATURE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2310,7 +2310,7 @@
                 PSA_ERROR_BUFFER_TOO_SMALL );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2355,7 +2355,7 @@
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2396,7 +2396,7 @@
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2467,7 +2467,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2588,7 +2588,7 @@
     PSA_ASSERT( psa_mac_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2643,7 +2643,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2681,7 +2681,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2759,7 +2759,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2909,7 +2909,7 @@
     PSA_ASSERT( psa_cipher_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2969,7 +2969,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3037,7 +3037,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3108,7 +3108,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3169,7 +3169,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3251,7 +3251,7 @@
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3354,7 +3354,7 @@
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3430,7 +3430,7 @@
     psa_destroy_key( handle );
     mbedtls_free( output_data );
     mbedtls_free( output_data2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3480,7 +3480,7 @@
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3536,7 +3536,7 @@
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3601,7 +3601,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3646,7 +3646,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3713,7 +3713,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3745,7 +3745,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3781,7 +3781,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3850,7 +3850,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3915,7 +3915,7 @@
     psa_destroy_key( handle );
     mbedtls_free( output );
     mbedtls_free( output2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3977,7 +3977,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4039,7 +4039,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4107,7 +4107,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4156,7 +4156,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4298,7 +4298,7 @@
     mbedtls_free( output_buffer );
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4379,7 +4379,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4439,7 +4439,7 @@
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4518,7 +4518,7 @@
     psa_key_derivation_abort( &operation );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4566,7 +4566,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4603,7 +4603,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4663,7 +4663,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4722,7 +4722,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( actual_output );
 }
 /* END_CASE */
@@ -4772,7 +4772,7 @@
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( output );
     mbedtls_free( changed );
 }
@@ -4818,7 +4818,7 @@
 exit:
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4919,7 +4919,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( e_read_buffer );
     mbedtls_free( exported );
 }
@@ -5016,7 +5016,7 @@
     }
 
     /* Shutdown and restart */
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Check key slot still contains key data */
@@ -5058,6 +5058,6 @@
         psa_open_key( key_id, &handle );
     }
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_entropy.function b/tests/suites/test_suite_psa_crypto_entropy.function
index 8576c7d..cd1b81f 100644
--- a/tests/suites/test_suite_psa_crypto_entropy.function
+++ b/tests/suites/test_suite_psa_crypto_entropy.function
@@ -77,7 +77,7 @@
 exit:
     mbedtls_free( seed );
     remove_seed_file( );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -105,12 +105,12 @@
     PSA_ASSERT( status );
     status = psa_crypto_init( );
     PSA_ASSERT( status );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     /* The seed is written by nv_seed callback functions therefore the injection will fail */
     status = mbedtls_psa_inject_entropy( seed, sizeof( seed ) );
     TEST_EQUAL( status, PSA_ERROR_NOT_PERMITTED );
 exit:
     remove_seed_file( );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_hash.function b/tests/suites/test_suite_psa_crypto_hash.function
index 90636b9..e15f335 100644
--- a/tests/suites/test_suite_psa_crypto_hash.function
+++ b/tests/suites/test_suite_psa_crypto_hash.function
@@ -31,7 +31,7 @@
                     actual_hash, actual_hash_length );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -52,7 +52,7 @@
                                  expected_hash->len ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -95,6 +95,6 @@
     } while( len++ != input->len );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_init.function b/tests/suites/test_suite_psa_crypto_init.function
index 7913158..eaf1b8b 100644
--- a/tests/suites/test_suite_psa_crypto_init.function
+++ b/tests/suites/test_suite_psa_crypto_init.function
@@ -138,7 +138,7 @@
         PSA_ASSERT( status );
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
 }
 /* END_CASE */
@@ -150,9 +150,9 @@
     for( i = 0; i < count; i++ )
     {
         PSA_ASSERT( psa_crypto_init( ) );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -166,7 +166,7 @@
     {
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
     status = psa_generate_random( random, sizeof( random ) );
     TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
@@ -186,7 +186,7 @@
     {
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
     psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
     status = psa_import_key( &attributes, data, sizeof( data ), &handle );
@@ -212,7 +212,7 @@
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -252,7 +252,7 @@
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -280,6 +280,6 @@
 
 exit:
     mbedtls_free( seed );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.function b/tests/suites/test_suite_psa_crypto_persistent_key.function
index 7e98fae..e4ab163 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.function
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.function
@@ -110,7 +110,7 @@
 
 exit:
     mbedtls_free( data );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -137,7 +137,7 @@
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -152,7 +152,7 @@
     TEST_EQUAL( handle, 0 );
 
     /* Shutdown and restart */
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Create another key in the same slot */
@@ -162,7 +162,7 @@
                                 &handle ) );
 
 exit:
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -192,7 +192,7 @@
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -209,7 +209,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_persistent_key( key_id );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
 }
 /* END_CASE */
 
@@ -241,7 +241,7 @@
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -276,7 +276,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     mbedtls_free( exported );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index a7bb596..fde3b4d 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -103,7 +103,7 @@
             PSA_ASSERT( psa_destroy_key( handle ) );
             break;
         case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
+            PSA_DONE( );
             PSA_ASSERT( psa_crypto_init( ) );
             break;
     }
@@ -114,7 +114,7 @@
     TEST_EQUAL( psa_close_key( handle ), PSA_ERROR_INVALID_HANDLE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -180,7 +180,7 @@
             PSA_ASSERT( psa_destroy_key( handle ) );
             break;
         case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
+            PSA_DONE( );
             PSA_ASSERT( psa_crypto_init( ) );
             break;
     }
@@ -236,7 +236,7 @@
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
     mbedtls_free( reexported );
 }
@@ -303,7 +303,7 @@
                     reexported, reexported_length );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
 }
 /* END_CASE */
@@ -322,7 +322,7 @@
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -350,7 +350,7 @@
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
 #endif
@@ -428,7 +428,7 @@
      * sure that the material is still alive. */
     if( target_lifetime != PSA_KEY_LIFETIME_VOLATILE )
     {
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
         PSA_ASSERT( psa_crypto_init( ) );
         PSA_ASSERT( psa_open_key( target_id, &target_handle ) );
     }
@@ -464,7 +464,7 @@
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -567,7 +567,7 @@
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -609,7 +609,7 @@
     PSA_ASSERT( psa_close_key( handle1 ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -657,7 +657,7 @@
     PSA_ASSERT( psa_close_key( handles[i - 1] ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( handles );
 }
 /* END_CASE */