Add tests

executor/mock/jobrunner.go

@@ -41,6 +41,10 @@ type JobInput struct {
 	IgnoreLaunchGuard bool
 	// Batch sets batch mode on the task
 	Batch string
+	// CPUBursting sets the CPU bursting protobuf attribute
+	CPUBursting bool
+	// CPU sets the CPU count resource attribute
+	CPU *int64
 	// StopTimeoutSeconds is the duration we wait after SIGTERM for the container to exit
 	KillWaitSeconds uint32
 }
@@ -211,7 +215,12 @@ func (jobRunner *JobRunner) StartJob(jobInput *JobInput) *JobRunResponse {
 	if id := jobInput.ImageDigest; id != "" {
 		ci.ImageDigest = protobuf.String(id)
 	}
+
+	ci.AllowCpuBursting = protobuf.Bool(jobInput.CPUBursting)
 	cpu := int64(1)
+	if jobInput.CPU != nil {
+		cpu = *jobInput.CPU
+	}
 	memMiB := int64(400)
 	diskMiB := uint64(100)
 

executor/mock/standalone/standalone_test.go

@@ -107,6 +107,9 @@ func TestStandalone(t *testing.T) {
 		testNewEnvironmentLocationNegative,
 		testOldEnvironmentLocationPositive,
 		testOldEnvironmentLocationNegative,
+		testNoCPUBursting,
+		testCPUBursting,
+		testTwoCPUs,
 	}
 	for _, fun := range testFunctions {
 		fullName := runtime.FuncForPC(reflect.ValueOf(fun).Pointer()).Name()
@@ -671,3 +674,45 @@ func testOldEnvironmentLocationNegative(t *testing.T, jobID string) {
 		t.Fail()
 	}
 }
+
+func testNoCPUBursting(t *testing.T, jobID string) {
+	ji := &mock.JobInput{
+		ImageName: ubuntu.name,
+		Version:   ubuntu.tag,
+		// Make sure quota is set
+		Entrypoint: `/bin/bash -c 'cat /sys/fs/cgroup/cpuacct/cpu.cfs_quota_us|grep -v - -1'`,
+		JobID:      jobID,
+	}
+	if !mock.RunJobExpectingSuccess(ji) {
+		t.Fail()
+	}
+}
+
+func testCPUBursting(t *testing.T, jobID string) {
+	ji := &mock.JobInput{
+		ImageName: ubuntu.name,
+		Version:   ubuntu.tag,
+		// Make sure quota is set
+		Entrypoint:  `/bin/bash -c 'cat /sys/fs/cgroup/cpuacct/cpu.cfs_quota_us|grep - -1'`,
+		JobID:       jobID,
+		CPUBursting: true,
+	}
+	if !mock.RunJobExpectingSuccess(ji) {
+		t.Fail()
+	}
+}
+
+func testTwoCPUs(t *testing.T, jobID string) {
+	var cpuCount int64 = 2
+	ji := &mock.JobInput{
+		ImageName: ubuntu.name,
+		Version:   ubuntu.tag,
+		// Make sure quota is set
+		Entrypoint: `/bin/bash -c 'cat /sys/fs/cgroup/cpuacct/cpu.shares|grep 200'`,
+		JobID:      jobID,
+		CPU:        &cpuCount,
+	}
+	if !mock.RunJobExpectingSuccess(ji) {
+		t.Fail()
+	}
+}