ct/go/poller/poller_test.go

package main

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"testing"

	"cloud.google.com/go/datastore"
	expect "github.com/stretchr/testify/assert"
	assert "github.com/stretchr/testify/require"
	"go.skia.org/infra/ct/go/ct_autoscaler"
	"go.skia.org/infra/ct/go/ctfe/admin_tasks"
	"go.skia.org/infra/ct/go/ctfe/capture_skps"
	"go.skia.org/infra/ct/go/ctfe/chromium_builds"
	"go.skia.org/infra/ct/go/ctfe/chromium_perf"
	"go.skia.org/infra/ct/go/ctfe/lua_scripts"
	"go.skia.org/infra/ct/go/ctfe/metrics_analysis"
	"go.skia.org/infra/ct/go/ctfe/pixel_diff"
	"go.skia.org/infra/ct/go/ctfe/task_common"
	ctfeutil "go.skia.org/infra/ct/go/ctfe/util"
	"go.skia.org/infra/ct/go/frontend"
	"go.skia.org/infra/go/ds"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/testutils"
	skutil "go.skia.org/infra/go/util"
)

// CommonCols without TsStarted or TsCompleted set.
func pendingCommonCols(kind ds.Kind) task_common.CommonCols {
	datastoreKey := datastore.Key{
		ID:   42,
		Kind: string(kind),
	}
	return task_common.CommonCols{
		DatastoreKey: &datastoreKey,
		TsAdded:      20080726180513,
		Username:     "nobody@chromium.org",
	}
}

// Given a command generated by one of the Execute methods in main.go, extracts the value of the
// run_id command-line flag. If not found, signals a test failure.
func getRunId(t *testing.T, cmd *exec.Command) string {
	regexp := regexp.MustCompile("^--run_id=(.*)$")
	for _, arg := range cmd.Args {
		match := regexp.FindStringSubmatch(arg)
		if match != nil && len(match) >= 2 {
			return match[1]
		}
	}
	assert.Contains(t, strings.Join(cmd.Args, " "), "--run_id=")
	assert.FailNow(t, "getRunId is broken")
	return ""
}

// Checks that the contents of filepath are expected; otherwise signals a test failure.
func assertFileContents(t *testing.T, filepath, expected string) {
	actual, err := ioutil.ReadFile(filepath)
	assert.NoError(t, err)
	assert.Equal(t, expected, string(actual))
}

func pendingChromiumPerfTask() ChromiumPerfTask {
	return ChromiumPerfTask{
		DatastoreTask: chromium_perf.DatastoreTask{
			CommonCols:           pendingCommonCols(ds.CHROMIUM_PERF_TASKS),
			Benchmark:            "benchmark",
			Platform:             "Linux",
			PageSets:             "All",
			RepeatRuns:           1,
			BenchmarkArgs:        "benchmarkargs",
			BrowserArgsNoPatch:   "banp",
			BrowserArgsWithPatch: "bawp",
			Description:          "description",
			ChromiumPatchGSPath:  "patches/abc.patch",
			SkiaPatchGSPath:      "patches/xyz.patch",
			CCList:               []string{"superman@krypton.com"},
		},
	}
}

func TestChromiumPerfExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
		runId := getRunId(t, cmd)
		assertFileContents(t, filepath.Join(os.TempDir(), runId+".chromium.patch"),
			"patches/abc.patch\n")
		assertFileContents(t, filepath.Join(os.TempDir(), runId+".skia.patch"),
			"patches/xyz.patch\n")
		return nil
	})
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingChromiumPerfTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "run_chromium_perf_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--description=description")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org,superman@krypton.com")
	expect.Contains(t, cmd.Args, "--benchmark_name=benchmark")
	expect.Contains(t, cmd.Args, "--target_platform=Linux")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	expect.Contains(t, cmd.Args, "--repeat_benchmark=1")
	expect.Contains(t, cmd.Args, "--benchmark_extra_args=benchmarkargs")
	expect.Contains(t, cmd.Args, "--browser_extra_args_nopatch=banp")
	expect.Contains(t, cmd.Args, "--browser_extra_args_withpatch=bawp")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 5, getPatchCalls)
}

func pendingPixelDiffTask() PixelDiffTask {
	return PixelDiffTask{
		DatastoreTask: pixel_diff.DatastoreTask{
			CommonCols:           pendingCommonCols(ds.PIXEL_DIFF_TASKS),
			PageSets:             "All",
			BenchmarkArgs:        "benchmarkargs",
			BrowserArgsNoPatch:   "banp",
			BrowserArgsWithPatch: "bawp",
			Description:          "description",
			ChromiumPatchGSPath:  "patches/abc.patch",
			SkiaPatchGSPath:      "patches/xyz.patch",
		},
	}
}

func TestPixelDiffExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingPixelDiffTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "pixel_diff_on_workers", cmd.Name)
	expect.Equal(t, len(cmd.Args), 14)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--description=description")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	expect.Contains(t, cmd.Args, "--benchmark_extra_args=benchmarkargs")
	expect.Contains(t, cmd.Args, "--browser_extra_args_nopatch=banp")
	expect.Contains(t, cmd.Args, "--browser_extra_args_withpatch=bawp")
	expect.Contains(t, cmd.Args, "--logtostderr")
	expect.Contains(t, cmd.Args, "--local=false")
	expect.Contains(t, cmd.Args, "--run_on_gce="+strconv.FormatBool(task.RunsOnGCEWorkers()))
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 3, getPatchCalls)
}

func pendingMetricsAnalysisTask() MetricsAnalysisTask {
	return MetricsAnalysisTask{
		DatastoreTask: metrics_analysis.DatastoreTask{
			CommonCols:          pendingCommonCols(ds.METRICS_ANALYSIS_TASKS),
			MetricName:          "loadingMetric",
			AnalysisOutputLink:  "http://test/outputlink",
			BenchmarkArgs:       "benchmarkargs",
			ValueColumnName:     "pct_001",
			Description:         "description",
			ChromiumPatchGSPath: "patches/abc.patch",
			CatapultPatchGSPath: "patches/xyz.path",
			CCList:              []string{"superman@krypton.com", "jorel@krypton.com"},
			TaskPriority:        100,
		},
	}
}

func TestMetricsAnalysisExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingMetricsAnalysisTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "metrics_analysis_on_workers", cmd.Name)
	expect.Equal(t, len(cmd.Args), 14)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--description=description")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org,superman@krypton.com,jorel@krypton.com")
	expect.Contains(t, cmd.Args, "--metric_name=loadingMetric")
	expect.Contains(t, cmd.Args, "--analysis_output_link=http://test/outputlink")
	expect.Contains(t, cmd.Args, "--value_column_name=pct_001")
	expect.Contains(t, cmd.Args, "--benchmark_extra_args=benchmarkargs")
	expect.Contains(t, cmd.Args, "--task_priority=100")
	expect.Contains(t, cmd.Args, "--logtostderr")
	expect.Contains(t, cmd.Args, "--local=false")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 3, getPatchCalls)
}

func pendingCaptureSkpsTask() CaptureSkpsTask {
	return CaptureSkpsTask{
		DatastoreTask: capture_skps.DatastoreTask{
			CommonCols:  pendingCommonCols(ds.CAPTURE_SKPS_TASKS),
			PageSets:    "All",
			ChromiumRev: "c14d891d44f0afff64e56ed7c9702df1d807b1ee",
			SkiaRev:     "586101c79b0490b50623e76c71a5fd67d8d92b08",
			Description: "description",
		},
	}
}

func TestCaptureSkpsExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingCaptureSkpsTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "capture_skps_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--description=description")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	expect.Contains(t, cmd.Args, "--chromium_build=c14d891d44f0af-586101c79b0490")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func pendingLuaScriptTaskWithAggregator(ctx context.Context) LuaScriptTask {
	return LuaScriptTask{
		DatastoreTask: lua_scripts.DatastoreTask{
			CommonCols:          pendingCommonCols(ds.LUA_SCRIPT_TASKS),
			PageSets:            "All",
			ChromiumRev:         "c14d891d44f0afff64e56ed7c9702df1d807b1ee",
			SkiaRev:             "586101c79b0490b50623e76c71a5fd67d8d92b08",
			LuaScript:           `print("lualualua")`,
			LuaAggregatorScript: `print("aaallluuu")`,
			Description:         "description",
		},
	}
}

func TestLuaScriptExecuteWithAggregator(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)
	task := pendingLuaScriptTaskWithAggregator(ctx)
	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
		runId := getRunId(t, cmd)
		assertFileContents(t, filepath.Join(os.TempDir(), runId+".lua"),
			`print("lualualua")`)
		assertFileContents(t, filepath.Join(os.TempDir(), runId+".aggregator"),
			`print("aaallluuu")`)
		return nil
	})

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "run_lua_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--description=description")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	expect.Contains(t, cmd.Args, "--chromium_build=c14d891d44f0af-586101c79b0490")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func TestLuaScriptExecuteWithoutAggregator(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
		runId := getRunId(t, cmd)
		assertFileContents(t, filepath.Join(os.TempDir(), runId+".lua"),
			`print("lualualua")`)
		_, err := os.Stat(filepath.Join(os.TempDir(), runId+".aggregator"))
		expect.True(t, os.IsNotExist(err))
		return nil
	})
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := LuaScriptTask{
		DatastoreTask: lua_scripts.DatastoreTask{
			CommonCols:          pendingCommonCols(ds.LUA_SCRIPT_TASKS),
			PageSets:            "All",
			ChromiumRev:         "c14d891d44f0afff64e56ed7c9702df1d807b1ee",
			SkiaRev:             "586101c79b0490b50623e76c71a5fd67d8d92b08",
			LuaScript:           `print("lualualua")`,
			LuaAggregatorScript: "",
			Description:         "description",
		},
	}
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "run_lua_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	expect.Contains(t, cmd.Args, "--chromium_build=c14d891d44f0af-586101c79b0490")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func pendingChromiumBuildTask() ChromiumBuildTask {
	return ChromiumBuildTask{
		DatastoreTask: chromium_builds.DatastoreTask{
			CommonCols:    pendingCommonCols(ds.CHROMIUM_BUILD_TASKS),
			ChromiumRev:   "c14d891d44f0afff64e56ed7c9702df1d807b1ee",
			ChromiumRevTs: 20080726180513,
			SkiaRev:       "586101c79b0490b50623e76c71a5fd67d8d92b08",
		},
	}
}

func TestChromiumBuildExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingChromiumBuildTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "build_chromium", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args,
		"--chromium_hash=c14d891d44f0afff64e56ed7c9702df1d807b1ee")
	expect.Contains(t, cmd.Args,
		"--skia_hash=586101c79b0490b50623e76c71a5fd67d8d92b08")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func pendingRecreatePageSetsTask() RecreatePageSetsTask {
	return RecreatePageSetsTask{
		RecreatePageSetsDatastoreTask: admin_tasks.RecreatePageSetsDatastoreTask{
			CommonCols: pendingCommonCols(ds.RECREATE_PAGESETS_TASKS),
			PageSets:   "All",
		},
	}
}

func TestRecreatePageSetsExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingRecreatePageSetsTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "create_pagesets_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func pendingRecreateWebpageArchivesTask() RecreateWebpageArchivesTask {
	return RecreateWebpageArchivesTask{
		RecreateWebpageArchivesDatastoreTask: admin_tasks.RecreateWebpageArchivesDatastoreTask{
			CommonCols:  pendingCommonCols(ds.RECREATE_WEBPAGE_ARCHIVES_TASKS),
			PageSets:    "All",
			ChromiumRev: "c14d891d44f0afff64e56ed7c9702df1d807b1ee",
			SkiaRev:     "586101c79b0490b50623e76c71a5fd67d8d92b08",
		},
	}
}

func TestRecreateWebpageArchivesExecute(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockRun.Run)

	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	task := pendingRecreateWebpageArchivesTask()
	err := task.Execute(ctx, mockGetPatchFromStorageFunc)
	assert.NoError(t, err)
	assert.Len(t, mockRun.Commands(), 1)
	cmd := mockRun.Commands()[0]
	expect.Equal(t, "capture_archives_on_workers", cmd.Name)
	expect.Contains(t, cmd.Args, "--task_id=42")
	expect.Contains(t, cmd.Args, "--emails=nobody@chromium.org")
	expect.Contains(t, cmd.Args, "--pageset_type=All")
	runId := getRunId(t, cmd)
	expect.Contains(t, cmd.Args, "--run_id="+runId)
	expect.NotNil(t, cmd.Timeout)
	expect.Equal(t, 0, getPatchCalls)
}

func TestAsPollerTask(t *testing.T) {
	testutils.SmallTest(t)
	ctx := context.Background()
	expect.Nil(t, asPollerTask(ctx, nil))
	{
		taskStruct := pendingChromiumPerfTask()
		taskInterface := asPollerTask(ctx, &taskStruct.DatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*ChromiumPerfTask))
	}
	{
		taskStruct := pendingCaptureSkpsTask()
		taskInterface := asPollerTask(ctx, &taskStruct.DatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*CaptureSkpsTask))
	}
	{
		taskStruct := pendingLuaScriptTaskWithAggregator(ctx)
		taskInterface := asPollerTask(ctx, &taskStruct.DatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*LuaScriptTask))
	}
	{
		taskStruct := pendingChromiumBuildTask()
		taskInterface := asPollerTask(ctx, &taskStruct.DatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*ChromiumBuildTask))
	}
	{
		taskStruct := pendingRecreatePageSetsTask()
		taskInterface := asPollerTask(ctx, &taskStruct.RecreatePageSetsDatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*RecreatePageSetsTask))
	}
	{
		taskStruct := pendingRecreateWebpageArchivesTask()
		taskInterface := asPollerTask(ctx, &taskStruct.RecreateWebpageArchivesDatastoreTask)
		expect.Equal(t, taskStruct, *taskInterface.(*RecreateWebpageArchivesTask))
	}
}

// Test that updateWebappTaskSetFailed works.
func TestUpdateWebappTaskSetFailed(t *testing.T) {
	testutils.SmallTest(t)
	task := pendingRecreateWebpageArchivesTask()
	mockServer := frontend.MockServer{}
	defer frontend.CloseTestServer(frontend.InitTestServer(&mockServer))
	err := updateWebappTaskSetFailed(&task)
	assert.NoError(t, err)
	assert.Len(t, mockServer.UpdateTaskReqs(), 1)
	updateReq := mockServer.UpdateTaskReqs()[0]
	assert.Equal(t, "/"+ctfeutil.UPDATE_RECREATE_WEBPAGE_ARCHIVES_TASK_POST_URI, updateReq.Url)
	assert.NoError(t, updateReq.Error)
	assert.Equal(t, "", updateReq.Vars.TsStarted)
	assert.NotEqual(t, "", updateReq.Vars.TsCompleted)
	assert.True(t, updateReq.Vars.Failure)
	assert.Equal(t, int64(0), updateReq.Vars.RepeatAfterDays)
	assert.False(t, updateReq.Vars.ClearRepeatAfterDays)
	assert.Equal(t, int64(42), updateReq.Vars.Id)
}

// Test that updateWebappTaskSetFailed returns an error when the server response indicates an error.
func TestUpdateWebappTaskSetFailedFailure(t *testing.T) {
	testutils.SmallTest(t)
	errstr := "You must be at least this tall to ride this ride."
	task := pendingRecreateWebpageArchivesTask()
	reqCount := 0
	mockServer := func(w http.ResponseWriter, r *http.Request) {
		reqCount++
		assert.Equal(t, "/"+ctfeutil.UPDATE_RECREATE_WEBPAGE_ARCHIVES_TASK_POST_URI,
			r.URL.Path)
		defer skutil.Close(r.Body)
		httputils.ReportError(w, r, fmt.Errorf(errstr), errstr)
	}
	defer frontend.CloseTestServer(frontend.InitTestServer(http.HandlerFunc(mockServer)))
	err := updateWebappTaskSetFailed(&task)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), errstr)
	assert.Equal(t, 1, reqCount)
}

func TestPollAndExecOnce(t *testing.T) {
	testutils.SmallTest(t)
	mockExec := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockExec.Run)
	task := pendingRecreateWebpageArchivesTask()
	mockCTAutoscaler := &ct_autoscaler.MockCTAutoscaler{}
	mockServer := frontend.MockServer{}
	mockServer.SetCurrentTask(&task.RecreateWebpageArchivesDatastoreTask)
	defer frontend.CloseTestServer(frontend.InitTestServer(&mockServer))
	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	wg := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg.Wait()
	// Expect only one poll.
	expect.Equal(t, 1, mockServer.OldestPendingTaskReqCount())
	expect.Equal(t, 1, mockCTAutoscaler.RegisterGCETaskTimesCalled)
	expect.Equal(t, 1, mockCTAutoscaler.UnregisterGCETaskTimesCalled)
	expect.Equal(t, 0, getPatchCalls)
	// Expect one command: capture_archives_on_workers ...
	commands := mockExec.Commands()
	assert.Len(t, commands, 1)
	expect.Equal(t, "capture_archives_on_workers", commands[0].Name)
	// No updates expected. (capture_archives_on_workers would send updates if it were
	// executed.)
	expect.Empty(t, mockServer.UpdateTaskReqs())
}

func TestPollAndExecOnceMultipleTasks(t *testing.T) {
	testutils.SmallTest(t)
	mockExec := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockExec.Run)
	task1 := pendingRecreateWebpageArchivesTask()
	mockCTAutoscaler := &ct_autoscaler.MockCTAutoscaler{}
	mockServer := frontend.MockServer{}
	mockServer.SetCurrentTask(&task1.RecreateWebpageArchivesDatastoreTask)
	defer frontend.CloseTestServer(frontend.InitTestServer(&mockServer))
	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	// Poll frontend and execute the first task.
	wg1 := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg1.Wait() // Wait for task to return to make asserting commands deterministic.
	// Update current task.
	task2 := pendingChromiumPerfTask()
	mockServer.SetCurrentTask(&task2.DatastoreTask)
	// Poll frontend and execute the second task.
	wg2 := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg2.Wait() // Wait for task to return to make asserting commands deterministic.

	// Expect two pending task requests.
	expect.Equal(t, 2, mockServer.OldestPendingTaskReqCount())
	expect.Equal(t, 1, mockCTAutoscaler.RegisterGCETaskTimesCalled)
	expect.Equal(t, 1, mockCTAutoscaler.UnregisterGCETaskTimesCalled)
	// Expect two commands: capture_archives_on_workers ...; run_chromium_perf_on_workers ...
	commands := mockExec.Commands()
	assert.Len(t, commands, 2)
	expect.Equal(t, "capture_archives_on_workers", commands[0].Name)
	expect.Equal(t, "run_chromium_perf_on_workers", commands[1].Name)
	// No updates expected when commands succeed.
	expect.Empty(t, mockServer.UpdateTaskReqs())
	expect.Equal(t, 5, getPatchCalls)
}

func TestPollAndExecOnceError(t *testing.T) {
	testutils.SmallTest(t)
	mockRun := exec.MockRun{}
	commandCollector := exec.CommandCollector{}
	commandCollector.SetDelegateRun(mockRun.Run)
	ctx := exec.NewContext(context.Background(), commandCollector.Run)
	task := pendingRecreateWebpageArchivesTask()
	mockCTAutoscaler := &ct_autoscaler.MockCTAutoscaler{}
	mockServer := frontend.MockServer{}
	mockServer.SetCurrentTask(&task.RecreateWebpageArchivesDatastoreTask)
	defer frontend.CloseTestServer(frontend.InitTestServer(&mockServer))
	mockRun.AddRule("capture_archives_on_workers", fmt.Errorf("workers too lazy"))
	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	wg := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg.Wait()
	// Expect only one poll.
	expect.Equal(t, 1, mockServer.OldestPendingTaskReqCount())
	expect.Equal(t, 1, mockCTAutoscaler.RegisterGCETaskTimesCalled)
	expect.Equal(t, 1, mockCTAutoscaler.UnregisterGCETaskTimesCalled)
	// Expect one command: capture_archives_on_workers ...
	commands := commandCollector.Commands()
	assert.Len(t, commands, 1)
	expect.Equal(t, "capture_archives_on_workers", commands[0].Name)
	// Expect an update marking task failed when command fails to execute.
	assert.Len(t, mockServer.UpdateTaskReqs(), 1)
	updateReq := mockServer.UpdateTaskReqs()[0]
	assert.Equal(t, "/"+ctfeutil.UPDATE_RECREATE_WEBPAGE_ARCHIVES_TASK_POST_URI, updateReq.Url)
	assert.NoError(t, updateReq.Error)
	assert.Equal(t, "", updateReq.Vars.TsStarted)
	assert.NotEqual(t, "", updateReq.Vars.TsCompleted)
	assert.True(t, updateReq.Vars.Failure)
	assert.Equal(t, int64(0), updateReq.Vars.RepeatAfterDays)
	assert.False(t, updateReq.Vars.ClearRepeatAfterDays)
	assert.Equal(t, int64(42), updateReq.Vars.Id)
	assert.Equal(t, 0, getPatchCalls)
}

func TestPollAndExecOnceNoTasks(t *testing.T) {
	testutils.SmallTest(t)
	mockCTAutoscaler := &ct_autoscaler.MockCTAutoscaler{}
	mockServer := frontend.MockServer{}
	mockServer.SetCurrentTask(nil)
	defer frontend.CloseTestServer(frontend.InitTestServer(&mockServer))
	mockExec := exec.CommandCollector{}
	ctx := exec.NewContext(context.Background(), mockExec.Run)
	getPatchCalls := 0
	mockGetPatchFromStorageFunc := func(patchId string) (string, error) {
		getPatchCalls++
		return patchId, nil
	}

	// Poll frontend, no tasks.
	wg1 := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg2 := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	wg3 := pollAndExecOnce(ctx, mockCTAutoscaler, mockGetPatchFromStorageFunc)
	// Expect three polls.
	wg1.Wait()
	wg2.Wait()
	wg3.Wait()
	expect.Equal(t, 3, mockServer.OldestPendingTaskReqCount())
	expect.Equal(t, 0, mockCTAutoscaler.RegisterGCETaskTimesCalled)
	expect.Equal(t, 0, mockCTAutoscaler.UnregisterGCETaskTimesCalled)
	expect.Equal(t, 0, getPatchCalls)
	// Expect no commands.
	expect.Empty(t, mockExec.Commands())
	// No updates expected.
	expect.Empty(t, mockServer.UpdateTaskReqs())
}