diff --git a/.pylintrc b/.pylintrc
index f9c97d5..4a1b6e5 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -70,6 +70,17 @@
 # Don't diplay statistics. Just the facts.
 reports=no
 
+[STRING]
+# Complain about
+# ```
+# list_of_strings = [
+#    'foo' # <-- missing comma
+#    'bar',
+#    'corge',
+# ]
+# ```
+check-str-concat-over-line-jumps=yes
+
 [VARIABLES]
 # Allow unused variables if their name starts with an underscore.
 # [unused-argument]
diff --git a/framework b/framework
index a39ba59..1e7b5d5 160000
--- a/framework
+++ b/framework
@@ -1 +1 @@
-Subproject commit a39ba59344fd4f1d0ee267ca414b9420d5dca9f5
+Subproject commit 1e7b5d54d3823b65fd4755bcf60f9ca39cfcbca3
diff --git a/programs/Makefile b/programs/Makefile
index b264290..a043fe1 100644
--- a/programs/Makefile
+++ b/programs/Makefile
@@ -47,6 +47,7 @@
 	../tf-psa-crypto/programs/psa/key_ladder_demo \
 	../tf-psa-crypto/programs/psa/psa_constant_names \
 	../tf-psa-crypto/programs/psa/psa_hash \
+	../tf-psa-crypto/programs/test/which_aes \
 	ssl/dtls_client \
 	ssl/dtls_server \
 	ssl/mini_client \
@@ -179,6 +180,10 @@
 	echo "  CC    psa/psa_hash.c"
 	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) ../tf-psa-crypto/programs/psa/psa_hash.c    $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
 
+../tf-psa-crypto/programs/test/which_aes$(EXEXT): ../tf-psa-crypto/programs/test/which_aes.c $(DEP)
+	echo "  CC    test/which_aes.c"
+	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) ../tf-psa-crypto/programs/test/which_aes.c    $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
+
 ssl/dtls_client$(EXEXT): ssl/dtls_client.c $(DEP)
 	echo "  CC    ssl/dtls_client.c"
 	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) ssl/dtls_client.c  $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
diff --git a/programs/test/generate_cpp_dummy_build.sh b/programs/test/generate_cpp_dummy_build.sh
index d27c7ae..7b4f520 100755
--- a/programs/test/generate_cpp_dummy_build.sh
+++ b/programs/test/generate_cpp_dummy_build.sh
@@ -52,6 +52,12 @@
         esac
     done
 
+    if [ -d "tf-psa-crypto/include/mbedtls" ]; then
+        for header in tf-psa-crypto/include/mbedtls/*.h; do
+            echo "#include \"${header#tf-psa-crypto/include/}\""
+        done
+    fi
+
     for header in tf-psa-crypto/include/psa/*.h; do
         case ${header#tf-psa-crypto/include/} in
             psa/crypto_config.h) :;; # not meant for direct inclusion
diff --git a/tests/Makefile b/tests/Makefile
index 87a6ca1..45231cd 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -324,25 +324,6 @@
 	perl -p -e 's/^(# *(define|ifndef) +\w+_)H\b/$${1}ALT_H/' $< >$@
 
 # Generate test library
-
-# Perl code that is executed to transform each original line from a library
-# source file into the corresponding line in the test driver copy of the
-# library. Add a LIBTESTDRIVER1_/libtestdriver1_ to mbedtls_xxx and psa_xxx
-# symbols.
-define libtestdriver1_rewrite :=
-	s!^(\s*#\s*include\s*[\"<])mbedtls/build_info.h!$${1}libtestdriver1/include/mbedtls/build_info.h!; \
-	s!^(\s*#\s*include\s*[\"<])mbedtls/mbedtls_config.h!$${1}libtestdriver1/include/mbedtls/mbedtls_config.h!; \
-	s!^(\s*#\s*include\s*[\"<])mbedtls/config_adjust_x509.h!$${1}libtestdriver1/include/mbedtls/config_adjust_x509.h!; \
-	s!^(\s*#\s*include\s*[\"<])mbedtls/config_adjust_ssl.h!$${1}libtestdriver1/include/mbedtls/config_adjust_ssl.h!; \
-	s!^(\s*#\s*include\s*[\"<])mbedtls/check_config.h!$${1}libtestdriver1/include/mbedtls/check_config.h!; \
-	s!^(\s*#\s*include\s*[\"<])mbedtls/!$${1}libtestdriver1/tf-psa-crypto/drivers/builtin/include/mbedtls/!; \
-	s!^(\s*#\s*include\s*[\"<])psa/!$${1}libtestdriver1/tf-psa-crypto/include/psa/!; \
-	s!^(\s*#\s*include\s*[\"<])tf-psa-crypto/!$${1}libtestdriver1/tf-psa-crypto/include/tf-psa-crypto/!; \
-	next if /^\s*#\s*include/; \
-	s/\b(?=MBEDTLS_|PSA_|TF_PSA_CRYPTO_)/LIBTESTDRIVER1_/g; \
-	s/\b(?=mbedtls_|psa_|tf_psa_crypto_)/libtestdriver1_/g;
-endef
-
 libtestdriver1.a:
 	rm -Rf ./libtestdriver1
 	mkdir ./libtestdriver1
@@ -384,12 +365,12 @@
 	# Prefix MBEDTLS_* PSA_* symbols with LIBTESTDRIVER1_ as well as
 	# mbedtls_* psa_* symbols with libtestdriver1_ to avoid symbol clash
 	# when this test driver library is linked with the Mbed TLS library.
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/library/*.[ch]
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/include/*/*.h
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/tf-psa-crypto/core/*.[ch]
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/tf-psa-crypto/include/*/*.h
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/tf-psa-crypto/drivers/builtin/include/*/*.h
-	perl -pi -e '$(libtestdriver1_rewrite)' ./libtestdriver1/tf-psa-crypto/drivers/builtin/src/*.[ch]
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/library/*.[ch]
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/include/*/*.h
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/tf-psa-crypto/core/*.[ch]
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/tf-psa-crypto/include/*/*.h
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/tf-psa-crypto/drivers/builtin/include/*/*.h
+	perl -i ./scripts/libtestdriver1_rewrite.pl ./libtestdriver1/tf-psa-crypto/drivers/builtin/src/*.[ch]
 
 	$(MAKE) -C ./libtestdriver1/library CFLAGS="-I../../ $(CFLAGS)" LDFLAGS="$(LDFLAGS)" libmbedcrypto.a
 	cp ./libtestdriver1/library/libmbedcrypto.a ../library/libtestdriver1.a
diff --git a/tests/scripts/components-platform.sh b/tests/scripts/components-platform.sh
index abae283..25cfd41 100644
--- a/tests/scripts/components-platform.sh
+++ b/tests/scripts/components-platform.sh
@@ -120,15 +120,17 @@
     msg "AES tests, test intrinsics"
     make clean
     make CC=gcc CFLAGS='-Werror -Wall -Wextra -mpclmul -msse2 -maes'
-    # check that we built intrinsics - this should be used by default when supported by the compiler
-    ./programs/test/selftest aes | grep "AESNI code" | grep -q "intrinsics"
+    # check that the intrinsics implementation is in use - this should be used by default when
+    # supported by the compiler
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI INTRINSICS"
 
     # test the asm implementation
     msg "AES tests, test assembly"
     make clean
     make CC=gcc CFLAGS='-Werror -Wall -Wextra -mno-pclmul -mno-sse2 -mno-aes'
-    # check that we built assembly - this should be built if the compiler does not support intrinsics
-    ./programs/test/selftest aes | grep "AESNI code" | grep -q "assembly"
+    # check that the assembly implementation is in use - this should be used if the compiler
+    # does not support intrinsics
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI ASSEMBLY"
 
     # test the plain C implementation
     scripts/config.py unset MBEDTLS_AESNI_C
@@ -136,21 +138,22 @@
     msg "AES tests, plain C"
     make clean
     make CC=gcc CFLAGS='-O2 -Werror'
-    # check that there is no AESNI code present
-    ./programs/test/selftest aes | not grep -q "AESNI code"
-    not grep -q "AES note: using AESNI" ./programs/test/selftest
-    grep -q "AES note: built-in implementation." ./programs/test/selftest
+    # check that the plain C implementation is present and the AESNI one is not
+    grep -q mbedtls_internal_aes_encrypt ./tf-psa-crypto/drivers/builtin/src/aes.o
+    not grep -q mbedtls_aesni_crypt_ecb ./tf-psa-crypto/drivers/builtin/src/aesni.o
+    # check that the built-in software implementation is in use
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "SOFTWARE"
 
-    # test the intrinsics implementation
+    # test the AESNI implementation
     scripts/config.py set MBEDTLS_AESNI_C
     scripts/config.py set MBEDTLS_AES_USE_HARDWARE_ONLY
     msg "AES tests, test AESNI only"
     make clean
     make CC=gcc CFLAGS='-Werror -Wall -Wextra -mpclmul -msse2 -maes'
-    ./programs/test/selftest aes | grep -q "AES note: using AESNI"
-    ./programs/test/selftest aes | not grep -q "AES note: built-in implementation."
-    grep -q "AES note: using AESNI" ./programs/test/selftest
-    not grep -q "AES note: built-in implementation." ./programs/test/selftest
+    # check that the AESNI implementation is present and the plain C one is not
+    grep -q mbedtls_aesni_crypt_ecb ./tf-psa-crypto/drivers/builtin/src/aesni.o
+    not grep -q mbedtls_internal_aes_encrypt ./tf-psa-crypto/drivers/builtin/src/aes.o
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI"
 }
 
 support_test_aesni_m32 () {
@@ -172,21 +175,22 @@
     make clean
     make CC=gcc CFLAGS='-m32 -Werror -Wall -Wextra' LDFLAGS='-m32'
     # check that we built intrinsics - this should be used by default when supported by the compiler
-    ./programs/test/selftest aes | grep "AESNI code" | grep -q "intrinsics"
-    grep -q "AES note: using AESNI" ./programs/test/selftest
-    grep -q "AES note: built-in implementation." ./programs/test/selftest
-    grep -q mbedtls_aesni_has_support ./programs/test/selftest
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI INTRINSICS"
+    # check that both the AESNI and plain C implementations are present
+    grep -q mbedtls_aesni_crypt_ecb ./tf-psa-crypto/drivers/builtin/src/aesni.o
+    grep -q mbedtls_internal_aes_encrypt ./tf-psa-crypto/drivers/builtin/src/aes.o
+    grep -q mbedtls_aesni_has_support ./tf-psa-crypto/programs/test/which_aes
 
     scripts/config.py set MBEDTLS_AESNI_C
     scripts/config.py set MBEDTLS_AES_USE_HARDWARE_ONLY
     msg "AES tests, test AESNI only"
     make clean
     make CC=gcc CFLAGS='-m32 -Werror -Wall -Wextra -mpclmul -msse2 -maes' LDFLAGS='-m32'
-    ./programs/test/selftest aes | grep -q "AES note: using AESNI"
-    ./programs/test/selftest aes | not grep -q "AES note: built-in implementation."
-    grep -q "AES note: using AESNI" ./programs/test/selftest
-    not grep -q "AES note: built-in implementation." ./programs/test/selftest
-    not grep -q mbedtls_aesni_has_support ./programs/test/selftest
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI"
+    # check that the AESNI implementation is present and the plain C one is not
+    grep -q mbedtls_aesni_crypt_ecb ./tf-psa-crypto/drivers/builtin/src/aesni.o
+    not grep -q mbedtls_internal_aes_encrypt ./tf-psa-crypto/drivers/builtin/src/aes.o
+    not grep -q mbedtls_aesni_has_support ./tf-psa-crypto/programs/test/which_aes
 }
 
 support_test_aesni_m32_clang () {
@@ -205,10 +209,11 @@
     make clean
     make CC=clang CFLAGS='-m32 -Werror -Wall -Wextra' LDFLAGS='-m32'
     # check that we built intrinsics - this should be used by default when supported by the compiler
-    ./programs/test/selftest aes | grep "AESNI code" | grep -q "intrinsics"
-    grep -q "AES note: using AESNI" ./programs/test/selftest
-    grep -q "AES note: built-in implementation." ./programs/test/selftest
-    grep -q mbedtls_aesni_has_support ./programs/test/selftest
+    ./tf-psa-crypto/programs/test/which_aes | grep -q "AESNI INTRINSICS"
+    # check that both the AESNI and plain C implementations are present
+    grep -q mbedtls_aesni_crypt_ecb ./tf-psa-crypto/drivers/builtin/src/aesni.o
+    grep -q mbedtls_internal_aes_encrypt ./tf-psa-crypto/drivers/builtin/src/aes.o
+    grep -q mbedtls_aesni_has_support ./tf-psa-crypto/programs/test/which_aes
 }
 
 support_build_aes_armce () {
diff --git a/tests/scripts/libtestdriver1_rewrite.pl b/tests/scripts/libtestdriver1_rewrite.pl
new file mode 100755
index 0000000..202575d
--- /dev/null
+++ b/tests/scripts/libtestdriver1_rewrite.pl
@@ -0,0 +1,41 @@
+#!/usr/bin/env perl
+
+# Perl code that is executed to transform each original line from a library
+# source file into the corresponding line in the test driver copy of the
+# library. Add a LIBTESTDRIVER1_/libtestdriver1_ to mbedtls_xxx and psa_xxx
+# symbols.
+
+# Copyright The Mbed TLS Contributors
+# SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+
+use warnings;
+use File::Basename;
+
+my @public_files = map { basename($_) } glob("../tf-psa-crypto/include/mbedtls/*.h");
+
+my $public_files_regex = join('|', map { quotemeta($_) } @public_files);
+
+while (<>) {
+    s!^(\s*#\s*include\s*[\"<])mbedtls/build_info.h!${1}libtestdriver1/include/mbedtls/build_info.h!;
+    s!^(\s*#\s*include\s*[\"<])mbedtls/mbedtls_config.h!${1}libtestdriver1/include/mbedtls/mbedtls_config.h!;
+    s!^(\s*#\s*include\s*[\"<])mbedtls/config_adjust_x509.h!${1}libtestdriver1/include/mbedtls/config_adjust_x509.h!;
+    s!^(\s*#\s*include\s*[\"<])mbedtls/config_adjust_ssl.h!${1}libtestdriver1/include/mbedtls/config_adjust_ssl.h!;
+    s!^(\s*#\s*include\s*[\"<])mbedtls/check_config.h!${1}libtestdriver1/include/mbedtls/check_config.h!;
+    # Files in include/mbedtls and drivers/builtin/include/mbedtls are both
+    # included in files via #include mbedtls/<file>.h, so when expanding to the
+    # full path make sure that files in include/mbedtls are not expanded
+    # to driver/builtin/include/mbedtls.
+    if ( $public_files_regex ) {
+        s!^(\s*#\s*include\s*[\"<])mbedtls/($public_files_regex)!${1}libtestdriver1/tf-psa-crypto/include/mbedtls/${2}!;
+    }
+    s!^(\s*#\s*include\s*[\"<])mbedtls/!${1}libtestdriver1/tf-psa-crypto/drivers/builtin/include/mbedtls/!;
+    s!^(\s*#\s*include\s*[\"<])psa/!${1}libtestdriver1/tf-psa-crypto/include/psa/!;
+    s!^(\s*#\s*include\s*[\"<])tf-psa-crypto/!${1}libtestdriver1/tf-psa-crypto/include/tf-psa-crypto/!;
+    if (/^\s*#\s*include/) {
+        print;
+        next;
+    }
+    s/\b(?=MBEDTLS_|PSA_|TF_PSA_CRYPTO_)/LIBTESTDRIVER1_/g;
+    s/\b(?=mbedtls_|psa_|tf_psa_crypto_)/libtestdriver1_/g;
+    print;
+}
diff --git a/tests/src/test_helpers/ssl_helpers.c b/tests/src/test_helpers/ssl_helpers.c
index 1d03eaf..1eed8ab 100644
--- a/tests/src/test_helpers/ssl_helpers.c
+++ b/tests/src/test_helpers/ssl_helpers.c
@@ -611,6 +611,7 @@
 {
     int i = 0;
     int ret = -1;
+    int ok = 0;
     mbedtls_test_ssl_endpoint_certificate *cert = NULL;
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     mbedtls_svc_key_id_t key_slot = MBEDTLS_SVC_KEY_ID_INIT;
@@ -733,7 +734,13 @@
                                     cert->pkey);
     TEST_ASSERT(ret == 0);
 
+    ok = 1;
+
 exit:
+    if (ret == 0 && !ok) {
+        /* Exiting due to a test assertion that isn't ret == 0 */
+        ret = -1;
+    }
     if (ret != 0) {
         test_ssl_endpoint_certificate_free(ep);
     }
@@ -865,6 +872,7 @@
 
     if (MBEDTLS_SSL_IS_CLIENT == endpoint_type) {
         ret = mbedtls_ssl_set_hostname(&(ep->ssl), "localhost");
+        TEST_EQUAL(ret, 0);
     }
 
 #if defined(MBEDTLS_SSL_PROTO_DTLS) && defined(MBEDTLS_SSL_SRV_C)
@@ -901,7 +909,13 @@
     TEST_EQUAL(mbedtls_ssl_get_user_data_n(&ep->ssl), user_data_n);
     mbedtls_ssl_set_user_data_p(&ep->ssl, ep);
 
+    return 0;
+
 exit:
+    if (ret == 0) {
+        /* Exiting due to a test assertion that isn't ret == 0 */
+        ret = -1;
+    }
     return ret;
 }
 
@@ -2598,6 +2612,7 @@
     mbedtls_ssl_session *session)
 {
     int ret = -1;
+    int ok = 0;
     unsigned char buf[64];
     mbedtls_test_ssl_endpoint client_ep, server_ep;
 
@@ -2634,10 +2649,16 @@
     ret = mbedtls_ssl_get_session(&(client_ep.ssl), session);
     TEST_EQUAL(ret, 0);
 
+    ok = 1;
+
 exit:
     mbedtls_test_ssl_endpoint_free(&client_ep, NULL);
     mbedtls_test_ssl_endpoint_free(&server_ep, NULL);
 
+    if (ret == 0 && !ok) {
+        /* Exiting due to a test assertion that isn't ret == 0 */
+        ret = -1;
+    }
     return ret;
 }
 #endif /* MBEDTLS_SSL_CLI_C && MBEDTLS_SSL_SRV_C &&
diff --git a/tf-psa-crypto b/tf-psa-crypto
index 4a9f29b..f936d86 160000
--- a/tf-psa-crypto
+++ b/tf-psa-crypto
@@ -1 +1 @@
-Subproject commit 4a9f29b05c661bd874c75d80339fcce00adea4e0
+Subproject commit f936d86b2587eb4a961cac5b3b95b949ee056ee6