diff --git a/library/common.h b/library/common.h
index a2c8a1e..9e4b031 100644
--- a/library/common.h
+++ b/library/common.h
@@ -46,6 +46,19 @@
 #define MBEDTLS_STATIC_TESTABLE static
 #endif
 
+#if defined(MBEDTLS_TEST_HOOKS)
+extern void (*mbedtls_test_hook_test_fail)( const char * test, int line, const char * file );
+#define MBEDTLS_TEST_HOOK_TEST_ASSERT( TEST ) \
+       do { \
+            if( ( ! ( TEST ) ) && ( ( *mbedtls_test_hook_test_fail ) != NULL ) ) \
+            { \
+              ( *mbedtls_test_hook_test_fail )( #TEST, __LINE__, __FILE__ ); \
+            } \
+    } while( 0 )
+#else
+#define MBEDTLS_TEST_HOOK_TEST_ASSERT( TEST )
+#endif /* defined(MBEDTLS_TEST_HOOKS) */
+
 /** Allow library to access its structs' private members.
  *
  * Although structs defined in header files are publicly available,
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index f5a8793..8312c07 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1002,10 +1002,7 @@
 
     if( slot->lock_count != 1 )
     {
-#if defined(MBEDTLS_TEST_HOOKS)
-        if( *mbedtls_test_hook_assert_test != NULL )
-            ( *mbedtls_test_hook_assert_test )( slot->lock_count == 1, __FILE__, __LINE__ );
-#endif
+        MBEDTLS_TEST_HOOK_TEST_ASSERT( slot->lock_count == 1 );
         status = PSA_ERROR_CORRUPTION_DETECTED;
     }
 
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 4a0763f..cb10f6e 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -34,7 +34,6 @@
 
 #include <stdlib.h>
 #include <string.h>
-#include "mbedtls/error.h"
 #if defined(MBEDTLS_PLATFORM_C)
 #include "mbedtls/platform.h"
 #else
@@ -413,13 +412,7 @@
         return( PSA_SUCCESS );
     }
 
-    slot->lock_count = 1;
-
-#if defined(MBEDTLS_TEST_HOOKS)
-    if( *mbedtls_test_hook_assert_test != NULL )
-        ( *mbedtls_test_hook_assert_test )( slot->lock_count > 0, __FILE__, __LINE__  );
-#endif
-
+    MBEDTLS_TEST_HOOK_TEST_ASSERT( slot->lock_count > 0 );
     return( PSA_ERROR_CORRUPTION_DETECTED );
 }