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