feat: improve interrupt handling (#996)

* refactor: use global context where applicable * don't wait for the ticker to cancel a poll * refactor: simplify interrupt handling * cleanup * fix: tests
saucelabs · Jan 15, 2025 · ef05397 · ef05397
1 parent 9db536f
commit ef05397
Show file tree

Hide file tree

Showing 19 changed files with 121 additions and 223 deletions.
diff --git a/internal/http/rdcservice.go b/internal/http/rdcservice.go
@@ -267,10 +267,12 @@ func (c *RDCService) PollJob(ctx context.Context, id string, interval, timeout t
 
 	for {
 		select {
+		case <-ctx.Done():
+			return job.Job{ID: id}, ctx.Err()
 		case <-ticker.C:
 			j, err := c.Job(ctx, id, realDevice)
 			if err != nil {
-				return job.Job{}, err
+				return job.Job{ID: id}, err
 			}
 
 			if job.Done(j.Status) {
@@ -280,7 +282,7 @@ func (c *RDCService) PollJob(ctx context.Context, id string, interval, timeout t
 		case <-deathclock.C:
 			j, err := c.Job(ctx, id, realDevice)
 			if err != nil {
-				return job.Job{}, err
+				return job.Job{ID: id}, err
 			}
 			j.TimedOut = true
 			return j, nil

diff --git a/internal/http/rdcservice_test.go b/internal/http/rdcservice_test.go
@@ -208,21 +208,21 @@ func TestRDCService_PollJob(t *testing.T) {
 			name:         "job not found error from external API",
 			client:       rdc,
 			jobID:        "3",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "3"},
 			expectedErr:  ErrJobNotFound,
 		},
 		{
 			name:         "http status is not 200, but 401 from external API",
 			client:       rdc,
 			jobID:        "4",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "4"},
 			expectedErr:  errors.New("unexpected statusCode: 401"),
 		},
 		{
 			name:         "unexpected status code from external API",
 			client:       rdc,
 			jobID:        "333",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "333"},
 			expectedErr:  errors.New("internal server error"),
 		},
 		{

diff --git a/internal/http/resto.go b/internal/http/resto.go
@@ -127,10 +127,12 @@ func (c *Resto) PollJob(ctx context.Context, id string, interval, timeout time.D
 
 	for {
 		select {
+		case <-ctx.Done():
+			return job.Job{ID: id}, ctx.Err()
 		case <-ticker.C:
 			j, err := c.Job(ctx, id, realDevice)
 			if err != nil {
-				return job.Job{}, err
+				return job.Job{ID: id}, err
 			}
 
 			if job.Done(j.Status) {
@@ -139,7 +141,7 @@ func (c *Resto) PollJob(ctx context.Context, id string, interval, timeout time.D
 		case <-deathclock.C:
 			j, err := c.Job(ctx, id, realDevice)
 			if err != nil {
-				return job.Job{}, err
+				return job.Job{ID: id}, err
 			}
 			j.TimedOut = true
 			return j, nil

diff --git a/internal/http/resto_test.go b/internal/http/resto_test.go
@@ -226,21 +226,21 @@ func TestResto_GetJobStatus(t *testing.T) {
 			name:         "user not found error from external API",
 			client:       resto,
 			jobID:        "3",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "3"},
 			expectedErr:  ErrJobNotFound,
 		},
 		{
 			name:         "http status is not 200, but 401 from external API",
 			client:       resto,
 			jobID:        "4",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "4"},
 			expectedErr:  errors.New("job status request failed; unexpected response code:'401', msg:''"),
 		},
 		{
 			name:         "unexpected status code from external API",
 			client:       resto,
 			jobID:        "333",
-			expectedResp: job.Job{},
+			expectedResp: job.Job{ID: "333"},
 			expectedErr:  errors.New("internal server error"),
 		},
 		{

diff --git a/internal/job/job.go b/internal/job/job.go
@@ -115,5 +115,5 @@ type Service interface {
 
 	// DownloadArtifacts downloads artifacts from a Job. Returns a list of
 	// file paths.
-	DownloadArtifacts(job Job, isLastAttempt bool) []string
+	DownloadArtifacts(ctx context.Context, job Job, isLastAttempt bool) []string
 }
diff --git a/internal/mocks/job.go b/internal/mocks/job.go
@@ -15,7 +15,7 @@ type FakeJobService struct {
 	PollJobFn  func(ctx context.Context, id string, interval time.Duration, timeout time.Duration) (job.Job, error)
 
 	UploadAssetFn            func(ctx context.Context, jobID string, realDevice bool, fileName string, contentType string, content []byte) error
-	DownloadArtifactFn       func(job job.Job, isLastAttempt bool) []string
+	DownloadArtifactFn       func(ctx context.Context, job job.Job, isLastAttempt bool) []string
 	GetJobAssetFileNamesFn   func(ctx context.Context, jobID string) ([]string, error)
 	GetJobAssetFileContentFn func(ctx context.Context, jobID, fileName string) ([]byte, error)
 }
@@ -34,10 +34,8 @@ func (s *FakeJobService) StopJob(
 	return s.StopJobFn(ctx, jobID, realDevice)
 }
 
-func (s *FakeJobService) DownloadArtifacts(
-	job job.Job, isLastAttempt bool,
-) []string {
-	return s.DownloadArtifactFn(job, isLastAttempt)
+func (s *FakeJobService) DownloadArtifacts(ctx context.Context, job job.Job, isLastAttempt bool) []string {
+	return s.DownloadArtifactFn(ctx, job, isLastAttempt)
 }
 
 // Job mock function