test: support recording REST API requests in acceptance tests [ENG-9440]

internal/acceptance_tests/recorded_transport.go

@@ -26,29 +26,80 @@ type recordedTransport struct {
 }
 
 type recording struct {
-	Request  graphqlRequest  `json:"request"`
-	Response graphqlResponse `json:"response"`
+	Request  recordedRequest  `json:"request"`
+	Response recordedResponse `json:"response"`
 }
 
-type graphqlRequest struct {
-	Query     string         `json:"query"`
+// recordedRequest is a recorded GraphQL or REST request.
+type recordedRequest struct {
+	// GraphQL requests
+	Query     string         `json:"query,omitempty"`
 	Variables map[string]any `json:"variables,omitempty"`
+
+	// REST requests
+	Method string          `json:"method,omitempty"`
+	Path   string          `json:"path,omitempty"`
+	Body   json.RawMessage `json:"body,omitempty"`
 }
 
-func (g graphqlRequest) Key() string {
-	key := g.Query
-	if len(g.Variables) > 0 {
-		varsJSON, _ := json.Marshal(g.Variables)
-		key = fmt.Sprintf("%s:%s", key, string(varsJSON))
+func (r recordedRequest) IsGraphQL() bool {
+	return r.Query != ""
+}
+
+func (r recordedRequest) Key() string {
+	if r.IsGraphQL() {
+		key := r.Query
+		if len(r.Variables) > 0 {
+			varsJSON, _ := json.Marshal(r.Variables)
+			key = fmt.Sprintf("%s:%s", key, string(varsJSON))
+		}
+		return key
+	}
+
+	key := r.Method + " " + r.Path
+	if len(r.Body) > 0 {
+		// recordings are saved indented, compact for consistent keys
+		var body bytes.Buffer
+		if err := json.Compact(&body, r.Body); err == nil {
+			key = fmt.Sprintf("%s:%s", key, body.String())
+		}
 	}
 	return key
 }
 
-type graphqlResponse struct {
+// recordedResponse is a recorded GraphQL or REST response.
+type recordedResponse struct {
+	// GraphQL responses
 	Data   any `json:"data,omitempty"`
 	Errors []struct {
 		Message string `json:"message"`
 	} `json:"errors,omitempty"`
+
+	// REST responses
+	StatusCode int             `json:"status_code,omitempty"`
+	Body       json.RawMessage `json:"body,omitempty"`
+	BodyText   string          `json:"body_text,omitempty"` // non-JSON body
+}
+
+func (r recordedResponse) RESTBody() []byte {
+	if r.BodyText != "" {
+		return []byte(r.BodyText)
+	}
+	return r.Body
+}
+
+// graphqlRequest is used to detect GraphQL requests.
+type graphqlRequest struct {
+	Query     string         `json:"query"`
+	Variables map[string]any `json:"variables"`
+}
+
+// graphqlResponse is used to decode GraphQL responses.
+type graphqlResponse struct {
+	Data   any `json:"data"`
+	Errors []struct {
+		Message string `json:"message"`
+	} `json:"errors"`
 }
 
 func newRecordedTransport(t *testing.T, testName string, mode string, wrapped http.RoundTripper) *recordedTransport {
@@ -126,33 +177,55 @@ func (rt *recordedTransport) RoundTrip(req *http.Request) (*http.Response, error
 	// http.Client.Do to traceback.
 	emptyResp := &http.Response{}
 
-	// Read the request body
-	var gqlReq graphqlRequest
-	body, err := io.ReadAll(req.Body)
-	if err != nil {
+	var body []byte
+	if req.Body != nil {
+		var err error
+		body, err = io.ReadAll(req.Body)
 		req.Body.Close()
-		return emptyResp, fmt.Errorf("failed to read request body: %v", err)
+		if err != nil {
+			return emptyResp, fmt.Errorf("failed to read request body: %v", err)
+		}
 	}
-	if err := json.Unmarshal(body, &gqlReq); err != nil {
-		req.Body.Close()
-		return emptyResp, fmt.Errorf("failed to decode request body: %v", err)
+
+	recReq, err := newRecordedRequest(req, body)
+	if err != nil {
+		return emptyResp, err
 	}
 
 	if rt.mode == TestModeRecord {
-		return rt.processRecordRequest(req, body, gqlReq)
+		return rt.processRecordRequest(req, body, recReq)
 	}
-	return rt.processReplayRequest(gqlReq)
+	return rt.processReplayRequest(recReq)
 }
 
-func (rt *recordedTransport) processRecordRequest(req *http.Request, body []byte, gqlReq graphqlRequest) (*http.Response, error) {
+// newRecordedRequest returns a recordedRequest for a GraphQL or REST request.
+func newRecordedRequest(req *http.Request, body []byte) (recordedRequest, error) {
+	var gqlReq graphqlRequest
+	if len(body) > 0 && json.Unmarshal(body, &gqlReq) == nil && gqlReq.Query != "" {
+		return recordedRequest{Query: gqlReq.Query, Variables: gqlReq.Variables}, nil
+	}
+
+	recReq := recordedRequest{Method: req.Method, Path: req.URL.RequestURI()}
+	if len(body) > 0 {
+		if !json.Valid(body) {
+			return recordedRequest{}, fmt.Errorf("request body is not valid JSON: %s", body)
+		}
+		recReq.Body = body
+	}
+	return recReq, nil
+}
+
+func (rt *recordedTransport) processRecordRequest(req *http.Request, body []byte, recReq recordedRequest) (*http.Response, error) {
 	// on errors, return an empty response, as returning nil causes
 	// http.Client.Do to traceback.
 	emptyResp := &http.Response{}
 
-	rt.t.Logf("Recording request with query: %s - %v", gqlReq.Query, gqlReq.Variables)
+	rt.t.Logf("Recording request: %s", recReq.Key())
 
 	// Make the real request with the original body
-	req.Body = io.NopCloser(bytes.NewReader(body))
+	if body != nil {
+		req.Body = io.NopCloser(bytes.NewReader(body))
+	}
 	resp, err := rt.wrapped.RoundTrip(req)
 	if err != nil {
 		return emptyResp, err
@@ -167,33 +240,57 @@ func (rt *recordedTransport) processRecordRequest(req *http.Request, body []byte
 		return emptyResp, err
 	}
 
-	// Parse the response
-	var gqlResp graphqlResponse
-	if err := json.Unmarshal(respBody, &gqlResp); err != nil {
-		return emptyResp, fmt.Errorf("failed to decode response body: %v", err)
+	var recResp recordedResponse
+	if recReq.IsGraphQL() {
+		var gqlResp graphqlResponse
+		if err := json.Unmarshal(respBody, &gqlResp); err != nil {
+			return emptyResp, fmt.Errorf("failed to decode response body: %v", err)
+		}
+		recResp.Data = gqlResp.Data
+		recResp.Errors = gqlResp.Errors
+	} else {
+		recResp.StatusCode = resp.StatusCode
+		if json.Valid(respBody) {
+			recResp.Body = respBody
+		} else {
+			recResp.BodyText = string(respBody)
+		}
 	}
 
-	rt.addRecording(recording{Request: gqlReq, Response: gqlResp})
+	rt.addRecording(recording{Request: recReq, Response: recResp})
 
 	// Return the response with a new body reader
 	resp.Body = io.NopCloser(bytes.NewReader(respBody))
 	return resp, nil
 }
 
-func (rt *recordedTransport) processReplayRequest(gqlReq graphqlRequest) (*http.Response, error) {
+func (rt *recordedTransport) processReplayRequest(recReq recordedRequest) (*http.Response, error) {
 	// on errors, return an empty response, as returning nil causes
 	// http.Client.Do to traceback.
 	emptyResp := &http.Response{}
 
-	key := gqlReq.Key()
-	rt.t.Logf("Attempting to replay request with query: %s - %v", gqlReq.Query, gqlReq.Variables)
+	key := recReq.Key()
+	rt.t.Logf("Attempting to replay request: %s", key)
 	rec, err := rt.getRecording(key)
 	if err != nil {
 		return emptyResp, err
 	}
 
+	if !recReq.IsGraphQL() {
+		statusCode := rec.Response.StatusCode
+		if statusCode == 0 {
+			statusCode = http.StatusOK
+		}
+		return &http.Response{
+			StatusCode: statusCode,
+			Status:     fmt.Sprintf("%d %s", statusCode, http.StatusText(statusCode)),
+			Body:       io.NopCloser(bytes.NewReader(rec.Response.RESTBody())),
+			Header:     make(http.Header),
+		}, nil
+	}
+
 	// Check that the recording matches
-	if gqlReq.Query != rec.Request.Query || !reflect.DeepEqual(gqlReq.Variables, rec.Request.Variables) {
+	if recReq.Query != rec.Request.Query || !reflect.DeepEqual(recReq.Variables, rec.Request.Variables) {
 		return emptyResp, fmt.Errorf(`
 Request doesn't match expected one:
 -- query
@@ -202,7 +299,7 @@ expected: %s
 -- parameters
 got     : %v
 expected: %v`,
-			gqlReq.Query, rec.Request.Query, gqlReq.Variables, rec.Request.Variables)
+			recReq.Query, rec.Request.Query, recReq.Variables, rec.Request.Variables)
 	}
 
 	// Create response
@@ -235,5 +332,5 @@ func (rt *recordedTransport) getRecording(key string) (recording, error) {
 			// no recording found in channel, fall out to the error case
 		}
 	}
-	return recording{}, fmt.Errorf("no recording found for query: %s", key)
+	return recording{}, fmt.Errorf("no recording found for request: %s", key)
 }