From 424ef05b1ac600af2d9aeae4449c89940cc7ca94 Mon Sep 17 00:00:00 2001
From: Noah S-C <noah@sourcegraph.com>
Date: Thu, 14 Jul 2022 14:33:04 -0700
Subject: [PATCH] gqltest: allow writing of request and response bodies to a
 sink (#38841)

---
 dev/codeintel-qa/internal/graphql.go |  8 +++++--
 internal/gqltestutil/client.go       | 35 ++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/dev/codeintel-qa/internal/graphql.go b/dev/codeintel-qa/internal/graphql.go
index 156ca3a8f01..bc723349440 100644
--- a/dev/codeintel-qa/internal/graphql.go
+++ b/dev/codeintel-qa/internal/graphql.go
@@ -1,11 +1,15 @@
 package internal
 
-import "github.com/sourcegraph/sourcegraph/internal/gqltestutil"
+import (
+	"os"
+
+	"github.com/sourcegraph/sourcegraph/internal/gqltestutil"
+)
 
 var client *gqltestutil.Client
 
 func InitializeGraphQLClient() (err error) {
-	client, err = gqltestutil.NewClient(SourcegraphEndpoint)
+	client, err = gqltestutil.NewClient(SourcegraphEndpoint, os.Stderr, os.Stderr)
 	return err
 }
 
diff --git a/internal/gqltestutil/client.go b/internal/gqltestutil/client.go
index 636dc680886..cda185ff807 100644
--- a/internal/gqltestutil/client.go
+++ b/internal/gqltestutil/client.go
@@ -59,7 +59,7 @@ func SignIn(baseURL, email, password string) (*Client, error) {
 
 // authenticate initializes an authenticated client with given request body.
 func authenticate(baseURL, path string, body any) (*Client, error) {
-	client, err := NewClient(baseURL)
+	client, err := NewClient(baseURL, nil, nil)
 	if err != nil {
 		return nil, errors.Wrap(err, "new client")
 	}
@@ -97,12 +97,23 @@ type Client struct {
 	csrfCookie    *http.Cookie
 	sessionCookie *http.Cookie
 
-	userID string
+	userID         string
+	requestLogger  io.StringWriter
+	responseLogger io.StringWriter
 }
 
 // NewClient instantiates a new client by performing a GET request then obtains the
-// CSRF token and cookie from its response, if there is one (old versions of Sourcegraph only.)
-func NewClient(baseURL string) (*Client, error) {
+// CSRF token and cookie from its response, if there is one (old versions of Sourcegraph only).
+// If request- or responseLogger are provided, the request and response bodies, respectively,
+// will be written to them for any GraphQL requests only.
+func NewClient(baseURL string, requestLogger, responseLogger io.StringWriter) (*Client, error) {
+	if requestLogger == nil {
+		requestLogger = io.Discard.(io.StringWriter)
+	}
+	if responseLogger == nil {
+		responseLogger = io.Discard.(io.StringWriter)
+	}
+
 	resp, err := http.Get(baseURL)
 	if err != nil {
 		return nil, errors.Wrap(err, "get URL")
@@ -114,6 +125,8 @@ func NewClient(baseURL string) (*Client, error) {
 		return nil, errors.Wrap(err, "read GET body")
 	}
 
+	responseLogger.WriteString(string(p))
+
 	csrfToken := extractCSRFToken(string(p))
 	var csrfCookie *http.Cookie
 	for _, cookie := range resp.Cookies() {
@@ -124,9 +137,11 @@ func NewClient(baseURL string) (*Client, error) {
 	}
 
 	return &Client{
-		baseURL:    baseURL,
-		csrfToken:  csrfToken,
-		csrfCookie: csrfCookie,
+		baseURL:        baseURL,
+		csrfToken:      csrfToken,
+		csrfCookie:     csrfCookie,
+		requestLogger:  requestLogger,
+		responseLogger: responseLogger,
 	}, nil
 }
 
@@ -139,6 +154,8 @@ func (c *Client) authenticate(path string, body any) error {
 		return errors.Wrap(err, "marshal body")
 	}
 
+	c.requestLogger.WriteString(string(p))
+
 	req, err := http.NewRequest("POST", c.baseURL+path, bytes.NewReader(p))
 	if err != nil {
 		return errors.Wrap(err, "new request")
@@ -257,6 +274,8 @@ func (c *Client) GraphQL(token, query string, variables map[string]any, target a
 		name = matches[2]
 	}
 
+	c.requestLogger.WriteString(string(body))
+
 	req, err := http.NewRequest("POST", fmt.Sprintf("%s/.api/graphql?%s", c.baseURL, name), bytes.NewReader(body))
 	if err != nil {
 		return err
@@ -286,6 +305,8 @@ func (c *Client) GraphQL(token, query string, variables map[string]any, target a
 		return errors.Wrap(err, "read response body")
 	}
 
+	c.responseLogger.WriteString(string(body))
+
 	// Check if the response format should be JSON
 	if strings.Contains(resp.Header.Get("Content-Type"), "application/json") {
 		// Try and see unmarshalling to errors