diff --git a/tests/scripts/check-generated-files.sh b/tests/scripts/check-generated-files.sh
index 09c850a..56d4884 100755
--- a/tests/scripts/check-generated-files.sh
+++ b/tests/scripts/check-generated-files.sh
@@ -31,9 +31,16 @@
 }
 
 if in_mbedtls_repo; then
-    library_dir='library'
+    if [ -d tf-psa-crypto ]; then
+        crypto_core_dir='tf-psa-crypto/core'
+        builtin_drivers_dir='tf-psa-crypto/drivers/builtin/src'
+    else
+        crypto_core_dir='library'
+        builtin_drivers_dir='library'
+    fi
 elif in_tf_psa_crypto_repo; then
-    library_dir='core'
+    crypto_core_dir='core'
+    builtin_drivers_dir='drivers/builtin/src/'
 else
     echo "Must be run from Mbed TLS root or TF-PSA-Crypto root" >&2
     exit 1
@@ -133,13 +140,14 @@
 check framework/scripts/generate_ecp_tests.py $(framework/scripts/generate_ecp_tests.py --list)
 check framework/scripts/generate_psa_tests.py $(framework/scripts/generate_psa_tests.py --list)
 check framework/scripts/generate_test_keys.py tests/src/test_keys.h
-check scripts/generate_driver_wrappers.py $library_dir/psa_crypto_driver_wrappers.h $library_dir/psa_crypto_driver_wrappers_no_static.c
+check scripts/generate_driver_wrappers.py ${crypto_core_dir}/psa_crypto_driver_wrappers.h \
+                                          ${crypto_core_dir}/psa_crypto_driver_wrappers_no_static.c
 
 # Additional checks for Mbed TLS only
 if in_mbedtls_repo; then
-    check scripts/generate_errors.pl library/error.c
+    check scripts/generate_errors.pl ${builtin_drivers_dir}/error.c
     check scripts/generate_query_config.pl programs/test/query_config.c
-    check scripts/generate_features.pl library/version_features.c
+    check scripts/generate_features.pl ${builtin_drivers_dir}/version_features.c
     check scripts/generate_ssl_debug_helpers.py library/ssl_debug_helpers_generated.c
     check framework/scripts/generate_test_cert_macros.py tests/src/test_certs.h
     # generate_visualc_files enumerates source files (library/*.c). It doesn't
diff --git a/tests/scripts/test_psa_compliance.py b/tests/scripts/test_psa_compliance.py
index d4e4979..7c09afc 100755
--- a/tests/scripts/test_psa_compliance.py
+++ b/tests/scripts/test_psa_compliance.py
@@ -39,7 +39,14 @@
     in_tf_psa_crypto_repo = build_tree.looks_like_tf_psa_crypto_root(root_dir)
 
     crypto_name = build_tree.crypto_library_filename(root_dir)
-    library_subdir = build_tree.crypto_core_directory(root_dir, relative=True)
+
+    # Temporary, while the crypto library is still located in the library
+    # directory. This will not be the case anymore when it will be built by
+    # the TF-PSA-Crypto build system.
+    if in_tf_psa_crypto_repo:
+        library_subdir = build_tree.crypto_core_directory(root_dir, relative=True)
+    else:
+        library_subdir = 'library'
 
     crypto_lib_filename = (library_build_dir + '/' +
                            library_subdir + '/' +