diff --git a/tests/test_slurm/run.sh b/tests/test_slurm/run.sh
index 54edac8..b43af44 100755
--- a/tests/test_slurm/run.sh
+++ b/tests/test_slurm/run.sh
@@ -1,9 +1,15 @@
 #!/bin/bash
 
-set -e
+exit_fail() {
+    err=$?
+    echo "Slurm test failed! Output is in $workdir"
+    exit $err
+}
+
+trap exit_fail ERR
 
-workdir=$(mktemp -d -p ./ disbatch-test.XXXX)
-cp Tasks $workdir
+workdir=$(mktemp -d -p $PWD disbatch-test.XXXX)
+cp Tasks Tasks_failfast $workdir
 cd $workdir
 
 # Run the test
@@ -11,16 +17,16 @@ salloc -n 2 disBatch Tasks
 
 # Check that all 3 tasks ran,
 # which means A.txt, B.txt, and C.txt exist
-success=0
-[[ -f A.txt && -f B.txt && -f C.txt ]] || success=$?
+[[ -f A.txt && -f B.txt && -f C.txt ]]
 
-cd - > /dev/null
+rm -f A.txt B.txt C.txt
 
-if [[ $success -eq 0 ]]; then
-    echo "Slurm test passed."
-    rm -rf $workdir
-else
-    echo "Slurm test failed! Output is in $workdir"
-fi
+# disBatch is expected to exit with a non-zero exit code here
+salloc -n 2 disbatch --fail-fast Tasks_failfast || true
+
+# check that we failed fast and didn't run any more tasks
+[[ ! -f A.txt ]]
 
-exit $success
+trap - ERR
+echo "Slurm test passed."
+rm -rf $workdir
diff --git a/tests/test_ssh/run.sh b/tests/test_ssh/run.sh
index 6d332c2..dfeb089 100755
--- a/tests/test_ssh/run.sh
+++ b/tests/test_ssh/run.sh
@@ -1,8 +1,14 @@
 #!/bin/bash
 
-set -e
+exit_fail() {
+    err=$?
+    echo "SSH test failed! Output is in $workdir"
+    exit $err
+}
+
+trap exit_fail ERR
 
-workdir=$(mktemp -d -p ./ disbatch-test.XXXX)
+workdir=$(mktemp -d -p $PWD disbatch-test.XXXX)
 cp Tasks Tasks_failfast $workdir
 cd $workdir
 
@@ -12,20 +18,15 @@ disBatch -s localhost:2 Tasks
 # Check that all 3 tasks ran,
 # which means A.txt, B.txt, and C.txt exist
 [[ -f A.txt && -f B.txt && -f C.txt ]]
-success=$?
 
-rm A.txt B.txt C.txt
-disbatch -s localhost:2 --fail-fast Tasks_failfast
-[[ ! -f A.txt ]]
-success=$((success + $?))
+rm -f A.txt B.txt C.txt
 
-cd - > /dev/null
+# disBatch is expected to exit with a non-zero exit code here
+disbatch -s localhost:2 --fail-fast Tasks_failfast || true
 
-if [[ $success -eq 0 ]]; then
-    echo "SSH test passed."
-    rm -rf $workdir
-else
-    echo "SSH test failed! Output is in $workdir"
-fi
+# check that we failed fast and didn't run any more tasks
+[[ ! -f A.txt ]]
 
-exit $success
+trap - ERR
+echo "SSH test passed."
+rm -rf $workdir