diff --git a/job.go b/job.go index 0c156f5..a351c04 100644 --- a/job.go +++ b/job.go @@ -96,15 +96,8 @@ func NewJob() Job { ramMb := aurora.NewResource() diskMb := aurora.NewResource() - resources := make(map[ResourceType]*aurora.Resource) - resources[CPU] = numCpus - resources[RAM] = ramMb - resources[DISK] = diskMb - - taskConfig.Resources = make(map[*aurora.Resource]bool) - taskConfig.Resources[numCpus] = true - taskConfig.Resources[ramMb] = true - taskConfig.Resources[diskMb] = true + resources := map[ResourceType]*aurora.Resource{CPU: numCpus, RAM: ramMb, DISK: diskMb} + taskConfig.Resources = map[*aurora.Resource]bool{numCpus: true, ramMb: true, diskMb: true} numCpus.NumCpus = new(float64) ramMb.RamMb = new(int64) @@ -178,7 +171,7 @@ func (j *AuroraJob) Disk(disk int64) Job { } func (j *AuroraJob) GPU(gpus int64) Job { - if _,ok := j.resources[GPU]; !ok { + if _, ok := j.resources[GPU]; !ok { numGPUs := &aurora.Resource{NumGpus: new(int64)} j.resources[GPU] = numGPUs j.TaskConfig().Resources[numGPUs] = true diff --git a/realis_e2e_test.go b/realis_e2e_test.go index dcd8e3c..777f1ef 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -28,6 +28,7 @@ import ( "github.com/paypal/gorealis/response" "github.com/pkg/errors" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) var r realis.Realis @@ -338,10 +339,9 @@ pulseLoop: select { case <-ticker.C: - fmt.Println("sending PulseJobUpdate....") resp, err = r.PulseJobUpdate(result.GetKey()) - assert.NotNil(t, resp) + require.NotNil(t, resp, "received a nil response from Aurora") assert.Nil(t, err) respDetail, err := r.JobUpdateDetails(updateQ) @@ -372,14 +372,13 @@ pulseLoop: _, err := r.AbortJobUpdate(*updateDetails[0].GetUpdate().GetSummary().GetKey(), "") assert.NoError(t, err) _, err = r.KillJob(job.JobKey()) - assert.NoError(t, err) - assert.FailNow(t, "timed out during pulse update test") + require.NoError(t, err, "timed out during pulse update test") } } resp, err = r.KillJob(job.JobKey()) - assert.NoError(t, err) - assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) + assert.NoError(t, err) + assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) } // Test configuring an executor that doesn't exist for CreateJob API