diff --git a/ci/build_device_and_tests b/ci/build_device_and_tests
index 9d11268..63d3ce3 100755
--- a/ci/build_device_and_tests
+++ b/ci/build_device_and_tests
@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+set -euo pipefail
 
-build/soong/soong_ui.bash --make-mode build_test_suites || exit $?
-$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites $@ || exit $?
+build/soong/soong_ui.bash --make-mode build_test_suites
+$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites --device-build $@
diff --git a/ci/build_test_suites b/ci/build_test_suites
index 9d11268..74470a8 100755
--- a/ci/build_test_suites
+++ b/ci/build_test_suites
@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+set -euo pipefail
 
-build/soong/soong_ui.bash --make-mode build_test_suites || exit $?
-$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites $@ || exit $?
+build/soong/soong_ui.bash --make-mode build_test_suites
+$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites $@
diff --git a/ci/build_test_suites.py b/ci/build_test_suites.py
index 623d256..19f0a2c 100644
--- a/ci/build_test_suites.py
+++ b/ci/build_test_suites.py
@@ -30,7 +30,7 @@
 import test_discovery_agent
 
 
-REQUIRED_ENV_VARS = frozenset(['TARGET_PRODUCT', 'TARGET_RELEASE', 'TOP'])
+REQUIRED_ENV_VARS = frozenset(['TARGET_PRODUCT', 'TARGET_RELEASE', 'TOP', 'DIST_DIR'])
 SOONG_UI_EXE_REL_PATH = 'build/soong/soong_ui.bash'
 LOG_PATH = 'logs/build_test_suites.log'
 
@@ -197,6 +197,11 @@
   argparser.add_argument(
       'extra_targets', nargs='*', help='Extra test suites to build.'
   )
+  argparser.add_argument(
+      '--device-build',
+      action='store_true',
+      help='Flag to indicate running a device build.',
+  )
 
   return argparser.parse_args(argv)
 
diff --git a/ci/build_test_suites_test.py b/ci/build_test_suites_test.py
index 26f4316..503a5f4 100644
--- a/ci/build_test_suites_test.py
+++ b/ci/build_test_suites_test.py
@@ -78,6 +78,12 @@
     with self.assert_raises_word(build_test_suites.Error, 'TOP'):
       build_test_suites.main([])
 
+  def test_missing_dist_dir_env_var_raises(self):
+    del os.environ['DIST_DIR']
+
+    with self.assert_raises_word(build_test_suites.Error, 'DIST_DIR'):
+      build_test_suites.main([])
+
   def test_invalid_arg_raises(self):
     invalid_args = ['--invalid_arg']
 
@@ -121,6 +127,7 @@
         'TARGET_RELEASE': 'release',
         'TARGET_PRODUCT': 'product',
         'TOP': str(self.fake_top),
+        'DIST_DIR': str(self.fake_top.joinpath('dist')),
     })
 
     self.mock_subprocess_run.return_value = 0