diff --git a/script/build_package.sh b/script/build_package.sh
index a52b3f3..e15eaf4 100755
--- a/script/build_package.sh
+++ b/script/build_package.sh
@@ -268,6 +268,14 @@
 		echo "Primary UART set to UART$uart."
 		set_run_env "primary_uart" "$uart"
 	fi
+
+	# UART used by payload(such as tftf, Linux) may not be the same as the
+	# primary UART. Set a run environment variable to track the payload
+	# UART which is tracked to check if the test has finished sucessfully.
+	if upon "$set_payload_uart"; then
+		echo "Payload uses UART$uart."
+		set_run_env "payload_uart" "$uart"
+	fi
 }
 
 # Extract a FIP in $1 using fiptool
diff --git a/script/run_package.sh b/script/run_package.sh
index 55e0870..346c4d3 100755
--- a/script/run_package.sh
+++ b/script/run_package.sh
@@ -162,6 +162,9 @@
 	die "No model path set by package!"
 fi
 
+# Assume primary UART is used by test payload as well
+payload_uart="${payload_uart:-$primary_uart}"
+
 # Launch model with parameters
 model_out="$run_root/model_log.txt"
 run_sh="$run_root/run.sh"
@@ -360,12 +363,12 @@
 
 	if [ "$u" = "$primary_uart" ]; then
 		star="*"
-		uart_name="primary_uart"
 	else
 		star=" "
-		uart_name="uart$u"
 	fi
 
+	uart_name="uart$u"
+
 	# Launch expect after exporting required variables
 	(
 	if [ -f "run/uart$u/env" ]; then
@@ -409,8 +412,8 @@
 		break
 	fi
 
-	# We're done if the primary UART exits success
-	if [ -f "$pid_dir/primary_uart.success" ]; then
+	# We're done if the payload UART exits success
+	if [ -f "$pid_dir/uart$payload_uart.success" ]; then
 		break
 	fi
 done