From c339ee1940defb51c92c02813a9e050eb60f517c Mon Sep 17 00:00:00 2001
From: Jan Hancic <jan.hancic@gmail.com>
Date: Thu, 8 Feb 2024 14:04:20 +0000
Subject: [PATCH] Download all inputs when running plz test & plz cover (#3069)

This is a follow-up for #3066
---
 src/core/state.go     |  1 -
 src/test/test_step.go | 30 ++++++++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/core/state.go b/src/core/state.go
index 1c9948c7e1..c05fd5ed10 100644
--- a/src/core/state.go
+++ b/src/core/state.go
@@ -1301,7 +1301,6 @@ func (state *BuildState) GetPreloadedSubincludes() []BuildLabel {
 
 // DownloadInputsIfNeeded downloads all the inputs (or runtime files) for a target if we are building remotely.
 func (state *BuildState) DownloadInputsIfNeeded(target *BuildTarget, runtime bool) error {
-	// TODO(jan): Remove this function once `DownloadAllInputs` has been fully implemented across both build & test.
 	if state.RemoteClient != nil {
 		state.LogBuildResult(target, TargetBuilding, "Downloading inputs...")
 		for input := range state.IterInputs(target, runtime) {
diff --git a/src/test/test_step.go b/src/test/test_step.go
index 238ebfae24..2f88f3d89d 100644
--- a/src/test/test_step.go
+++ b/src/test/test_step.go
@@ -69,16 +69,7 @@ func test(state *core.BuildState, label core.BuildLabel, target *core.BuildTarge
 
 	// If the user passed --shell then just prepare the directory.
 	if state.PrepareOnly {
-		if err := state.DownloadInputsIfNeeded(target, true); err != nil {
-			state.LogBuildError(label, core.TargetTestFailed, err, "Failed to download test inputs")
-			return
-		}
-		if err := core.PrepareRuntimeDir(state, target, target.TestDir(run)); err != nil {
-			state.LogBuildError(label, core.TargetTestFailed, err, "Failed to prepare test directory")
-			return
-		}
-		target.SetState(core.Stopped)
-		state.LogBuildResult(target, core.TargetTestStopped, "Test stopped")
+		prepareOnly(state, label, target, run)
 		return
 	}
 
@@ -268,6 +259,25 @@ func doFlakeRun(state *core.BuildState, target *core.BuildTarget, run int, runRe
 	return results, coverage
 }
 
+func prepareOnly(state *core.BuildState, label core.BuildLabel, target *core.BuildTarget, runNumber int) {
+	if state.RemoteClient != nil {
+		// Targets were built remotely so we can simply download the inputs and place them in the
+		// tmp/ folder and exit.
+		if err := state.DownloadAllInputs(target, target.TestDir(runNumber), true); err != nil {
+			state.LogBuildError(label, core.TargetTestFailed, err, "Failed to download test inputs")
+			return
+		}
+	} else {
+		if err := core.PrepareRuntimeDir(state, target, target.TestDir(runNumber)); err != nil {
+			state.LogBuildError(label, core.TargetTestFailed, err, "Failed to prepare test directory")
+			return
+		}
+	}
+
+	target.SetState(core.Stopped)
+	state.LogBuildResult(target, core.TargetTestStopped, "Test stopped")
+}
+
 func getFlakeStatus(flake int, flakes int) string {
 	if flakes == 1 {
 		return "Testing..."