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/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 () {