Refactor agent integration test API

agent/integration/config_allowlisting_integration_test.go

@@ -126,19 +126,22 @@ func TestConfigAllowlisting(t *testing.T) {
 			}
 
 			e := createTestAgentEndpoint()
-			server := e.server(jobID)
+			server := e.server()
 			defer server.Close()
 
 			mb := mockBootstrap(t)
 			tc.mockBootstrapExpectation(mb)
 			defer mb.CheckAndClose(t)
 
-			runJob(t, context.Background(), testRunJobConfig{
+			err := runJob(t, context.Background(), testRunJobConfig{
 				job:           job,
 				server:        server,
 				agentCfg:      tc.agentConfig,
 				mockBootstrap: mb,
 			})
+			if err != nil {
+				t.Fatalf("runJob() error = %v", err)
+			}
 
 			finishedJob := e.finishesFor(t, jobID)[0]
 

agent/integration/job_environment_integration_test.go

@@ -36,13 +36,16 @@ func TestWhenCachePathsSetInJobStep_CachePathsEnvVarIsSet(t *testing.T) {
 
 	// create a mock agent API
 	e := createTestAgentEndpoint()
-	server := e.server("my-job-id")
+	server := e.server()
 	defer server.Close()
 
-	runJob(t, ctx, testRunJobConfig{
+	err := runJob(t, ctx, testRunJobConfig{
 		job:           job,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{},
 		mockBootstrap: mb,
 	})
+	if err != nil {
+		t.Fatalf("runJob() error = %v", err)
+	}
 }

agent/integration/job_runner_integration_test.go

@@ -94,7 +94,7 @@ func TestPreBootstrapHookScripts(t *testing.T) {
 
 			// Creates a mock agent API
 			e := createTestAgentEndpoint()
-			server := e.server(defaultJobID)
+			server := e.server()
 			t.Cleanup(server.Close)
 
 			j := &api.Job{
@@ -111,12 +111,15 @@ func TestPreBootstrapHookScripts(t *testing.T) {
 			} else {
 				mb.Expect().NotCalled()
 			}
-			runJob(t, ctx, testRunJobConfig{
+			err = runJob(t, ctx, testRunJobConfig{
 				job:           j,
 				server:        server,
 				agentCfg:      agent.AgentConfiguration{HooksPath: hooksDir},
 				mockBootstrap: mb,
 			})
+			if err != nil {
+				t.Fatalf("runJob() error = %v", err)
+			}
 
 			mb.CheckAndClose(t)
 		})
@@ -151,19 +154,22 @@ func TestPreBootstrapHookRefusesJob(t *testing.T) {
 
 	// create a mock agent API
 	e := createTestAgentEndpoint()
-	server := e.server("my-job-id")
+	server := e.server()
 	defer server.Close()
 
 	mb := mockBootstrap(t)
 	mb.Expect().NotCalled() // The bootstrap won't be called, as the pre-bootstrap hook failed
 	defer mb.CheckAndClose(t)
 
-	runJob(t, ctx, testRunJobConfig{
+	err = runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{HooksPath: hooksDir},
 		mockBootstrap: mb,
 	})
+	if err != nil {
+		t.Fatalf("runJob() error = %v", err)
+	}
 
 	job := e.finishesFor(t, jobID)[0]
 
@@ -199,15 +205,19 @@ func TestJobRunner_WhenBootstrapExits_ItSendsTheExitStatusToTheAPI(t *testing.T)
 			mb.Expect().Once().AndExitWith(exit)
 
 			e := createTestAgentEndpoint()
-			server := e.server("my-job-id")
+			server := e.server()
 			defer server.Close()
 
-			runJob(t, ctx, testRunJobConfig{
+			err := runJob(t, ctx, testRunJobConfig{
 				job:           j,
 				server:        server,
 				agentCfg:      agent.AgentConfiguration{},
 				mockBootstrap: mb,
 			})
+			if err != nil {
+				t.Fatalf("runJob() error = %v", err)
+			}
+
 			finish := e.finishesFor(t, "my-job-id")[0]
 
 			if got, want := finish.ExitStatus, strconv.Itoa(exit); got != want {
@@ -244,15 +254,18 @@ func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
 
 	// create a mock agent API
 	e := createTestAgentEndpoint()
-	server := e.server("my-job-id")
+	server := e.server()
 	defer server.Close()
 
-	runJob(t, ctx, testRunJobConfig{
+	err := runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{},
 		mockBootstrap: mb,
 	})
+	if err != nil {
+		t.Fatalf("runJob() error = %v", err)
+	}
 }
 
 // TODO 2023-07-17: What is this testing? How is it testing it?
@@ -281,15 +294,18 @@ func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
 
 	// create a mock agent API
 	e := createTestAgentEndpoint()
-	server := e.server("my-job-id")
+	server := e.server()
 	defer server.Close()
 
-	runJob(t, ctx, testRunJobConfig{
+	err := runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{},
 		mockBootstrap: mb,
 	})
+	if err != nil {
+		t.Fatalf("runJob() error = %v", err)
+	}
 }
 
 func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
@@ -317,7 +333,7 @@ func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
 
 	// create a mock agent API
 	e := createTestAgentEndpoint()
-	server := e.server("my-job-id")
+	server := e.server()
 	defer server.Close()
 
 	runJob(t, ctx, testRunJobConfig{

agent/integration/job_verification_integration_test.go

@@ -556,7 +556,7 @@ func TestJobVerification(t *testing.T) {
 
 			// create a mock agent API
 			e := createTestAgentEndpoint()
-			server := e.server(tc.job.ID)
+			server := e.server()
 			defer server.Close()
 
 			mb := mockBootstrap(t)
@@ -569,13 +569,16 @@ func TestJobVerification(t *testing.T) {
 			}
 
 			tc.job.Step = signStep(t, tc.signingKey, pipelineUploadEnv, stepWithInvariants)
-			runJob(t, ctx, testRunJobConfig{
+			err := runJob(t, ctx, testRunJobConfig{
 				job:              &tc.job,
 				server:           server,
 				agentCfg:         tc.agentConf,
 				mockBootstrap:    mb,
 				verificationJWKS: tc.verificationJWKS,
 			})
+			if err != nil {
+				t.Fatalf("runJob() error = %v", err)
+			}
 
 			job := e.finishesFor(t, tc.job.ID)[0]
 

agent/integration/test_helpers.go

@@ -47,7 +47,7 @@ type testRunJobConfig struct {
 	verificationJWKS jwk.Set
 }
 
-func runJob(t *testing.T, ctx context.Context, cfg testRunJobConfig) {
+func runJob(t *testing.T, ctx context.Context, cfg testRunJobConfig) error {
 	t.Helper()
 
 	l := logger.Discard
@@ -70,13 +70,16 @@ func runJob(t *testing.T, ctx context.Context, cfg testRunJobConfig) {
 		AgentConfiguration: cfg.agentCfg,
 		MetricsScope:       scope,
 	})
+
 	if err != nil {
 		t.Fatalf("agent.NewJobRunner() error = %v", err)
 	}
 
 	if err := jr.Run(context.Background()); err != nil {
-		t.Errorf("jr.Run() = %v", err)
+		return err
 	}
+
+	return nil
 }
 
 type testAgentEndpoint struct {
@@ -125,33 +128,116 @@ func (tae *testAgentEndpoint) logsFor(t *testing.T, jobID string) string {
 	return strings.Join(logChunks, "")
 }
 
-func (t *testAgentEndpoint) server(jobID string) *httptest.Server {
-	return httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-		t.mtx.Lock()
-		defer t.mtx.Unlock()
-
-		b, _ := io.ReadAll(req.Body)
-		t.calls[req.URL.Path] = append(t.calls[req.URL.Path], b)
-
-		switch req.URL.Path {
-		case "/jobs/" + jobID:
-			rw.WriteHeader(http.StatusOK)
-			fmt.Fprintf(rw, `{"state":"running"}`)
-		case "/jobs/" + jobID + "/start":
-			rw.WriteHeader(http.StatusOK)
-		case "/jobs/" + jobID + "/chunks":
-			sequence := req.URL.Query().Get("sequence")
-			seqNo, _ := strconv.Atoi(sequence)
-			r, _ := gzip.NewReader(bytes.NewBuffer(b))
-			uz, _ := io.ReadAll(r)
-			t.logChunks[seqNo] = string(uz)
-			rw.WriteHeader(http.StatusCreated)
-		case "/jobs/" + jobID + "/finish":
-			rw.WriteHeader(http.StatusOK)
-		default:
-			http.Error(rw, fmt.Sprintf("not found; method = %q, path = %q", req.Method, req.URL.Path), http.StatusNotFound)
+type route struct {
+	Path   string
+	Method string
+	http.HandlerFunc
+}
+
+func (t *testAgentEndpoint) getJobsHandler() http.HandlerFunc {
+	return func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+		fmt.Fprintf(rw, `{"state":"running"}`)
+	}
+}
+
+func (t *testAgentEndpoint) chunksHandler() http.HandlerFunc {
+	return func(rw http.ResponseWriter, req *http.Request) {
+		b, err := io.ReadAll(req.Body)
+		if err != nil {
+			rw.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
+		sequence := req.URL.Query().Get("sequence")
+		seqNo, err := strconv.Atoi(sequence)
+		if err != nil {
+			rw.WriteHeader(http.StatusBadRequest)
+			return
+		}
+
+		r, err := gzip.NewReader(bytes.NewBuffer(b))
+		if err != nil {
+			rw.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
+		uz, err := io.ReadAll(r)
+		if err != nil {
+			rw.WriteHeader(http.StatusInternalServerError)
+			return
 		}
-	}))
+
+		t.logChunks[seqNo] = string(uz)
+		rw.WriteHeader(http.StatusCreated)
+	}
+}
+
+func (t *testAgentEndpoint) defaultRoutes() []route {
+	return []route{
+		{
+			Method:      "GET",
+			Path:        "/jobs/",
+			HandlerFunc: t.getJobsHandler(),
+		},
+		{
+			Method:      "PUT",
+			Path:        "/jobs/{id}/start",
+			HandlerFunc: func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) },
+		},
+		{
+			Method:      "POST",
+			Path:        "/jobs/{id}/chunks",
+			HandlerFunc: t.chunksHandler(),
+		},
+		{
+			Method:      "PUT",
+			Path:        "/jobs/{id}/finish",
+			HandlerFunc: func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) },
+		},
+	}
+}
+
+func (t *testAgentEndpoint) server(extraRoutes ...route) *httptest.Server {
+	mux := http.NewServeMux()
+
+	defaultRoutes := t.defaultRoutes()
+	routesUniq := make(map[string]http.HandlerFunc, len(defaultRoutes))
+	for _, r := range defaultRoutes {
+		routesUniq[fmt.Sprintf("%s %s", r.Method, r.Path)] = r.HandlerFunc
+	}
+
+	// extra routes overwrite default routes if they conflict
+	for _, r := range extraRoutes {
+		routesUniq[fmt.Sprintf("%s %s", r.Method, r.Path)] = r.HandlerFunc
+	}
+
+	wrapRecordRequest := func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			// wait a minute, what's going on here?
+			// well, we want to read the body of the request, but because HTTP response bodies are io.ReadClosers, they can only
+			// be read once. So we read the body, then write it back into the request body so that the next handler can read it.
+			b, _ := io.ReadAll(req.Body)
+			req.Body.Close()
+
+			// bytes.NewBuffer takes ownership of the slice, so we need to copy it
+			newBodyBytes := make([]byte, len(b))
+			copy(newBodyBytes, b)
+			req.Body = io.NopCloser(bytes.NewBuffer(newBodyBytes))
+
+			t.mtx.Lock()
+			t.calls[req.URL.Path] = append(t.calls[req.URL.Path], b)
+			t.mtx.Unlock()
+
+			next.ServeHTTP(rw, req)
+		})
+	}
+
+	for path, handler := range routesUniq {
+		mux.Handle(path, wrapRecordRequest(handler))
+	}
+
+	return httptest.NewServer(mux)
 }
 
 func mockPreBootstrap(t *testing.T, hooksDir string) *bintest.Mock {