Download all inputs when running plz test & plz cover (#3069)

This is a follow-up for #3066

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) {

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..."