diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 38a60b4..860e071 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -32,6 +32,18 @@
 #include "mbedtls/psa_util.h"
 #endif
 
+#if defined(MBEDTLS_MD_LIGHT)
+#include "mbedtls/md.h"
+#endif
+
+#if defined(MBEDTLS_MD_SOME_PSA)
+#define MD_PSA_INIT()   PSA_INIT()
+#define MD_PSA_DONE()   PSA_DONE()
+#else /* MBEDTLS_MD_SOME_PSA */
+#define MD_PSA_INIT() ((void) 0)
+#define MD_PSA_DONE() ((void) 0)
+#endif /* MBEDTLS_MD_SOME_PSA */
+
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 /** Initialize the PSA Crypto subsystem. */
 #define PSA_INIT() PSA_ASSERT(psa_crypto_init())
diff --git a/tests/suites/test_suite_md.function b/tests/suites/test_suite_md.function
index 32c9c8c..4438fa1 100644
--- a/tests/suites/test_suite_md.function
+++ b/tests/suites/test_suite_md.function
@@ -1,13 +1,5 @@
 /* BEGIN_HEADER */
 #include "mbedtls/md.h"
-
-#if defined(MBEDTLS_MD_SOME_PSA)
-#define MD_PSA_INIT()   PSA_INIT()
-#define MD_PSA_DONE()   PSA_DONE()
-#else /* MBEDTLS_MD_SOME_PSA */
-#define MD_PSA_INIT() ((void) 0)
-#define MD_PSA_DONE() ((void) 0)
-#endif /* MBEDTLS_MD_SOME_PSA */
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES