Build: psa-arch-test: Add toolchain selection

Change-Id: If5192a2080aadfb07964581fe41951bb8b780e3a
Signed-off-by: Kevin Peng <kevin.peng@arm.com>
diff --git a/tests_psa_arch/CMakeLists.txt b/tests_psa_arch/CMakeLists.txt
index dc4f2e2..2dfb91f 100644
--- a/tests_psa_arch/CMakeLists.txt
+++ b/tests_psa_arch/CMakeLists.txt
@@ -60,6 +60,16 @@
     set(TARGET tgt_dev_apis_tfm_${PSA_API_TEST_TARGET})
 endif()
 
+# Toolchain
+if (NOT TOOLCHAIN)
+    if(TFM_TOOLCHAIN_FILE MATCHES ".*toolchain_ns_(.+).cmake")
+        set(TOOLCHAIN ${CMAKE_MATCH_1})
+    else()
+        # TFM_TOOLCHAIN_FILE might not set, use the default one.
+        set(TOOLCHAIN GNUARM)
+    endif()
+endif()
+
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_SOURCE_DIR}/../cmake)
 include(toolchain_selection)
 
diff --git a/tests_psa_arch/spe/config/config_test_psa_api.cmake b/tests_psa_arch/spe/config/config_test_psa_api.cmake
index d46a912..e6c99b6 100644
--- a/tests_psa_arch/spe/config/config_test_psa_api.cmake
+++ b/tests_psa_arch/spe/config/config_test_psa_api.cmake
@@ -39,3 +39,12 @@
 
 # Set default value for INCLUDE_PANIC_TESTS explicitly
 set(INCLUDE_PANIC_TESTS     0   CACHE BOOL      "Include panic tests")
+
+if (NOT TOOLCHAIN)
+    if(TFM_TOOLCHAIN_FILE MATCHES ".*toolchain_(.+).cmake")
+        set(TOOLCHAIN ${CMAKE_MATCH_1})
+    else()
+        # TFM_TOOLCHAIN_FILE might not set, use the default one.
+        set(TOOLCHAIN GNUARM)
+    endif()
+endif()