From 787dbe909912ab5c4894cb21218cb39d6626e95a Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Fri, 31 Jan 2025 19:01:23 +0530
Subject: [PATCH] Add request body assertions to acceptance tests (#2263)

## Changes
With this PR, any acceptance tests that define custom server stubs in
`test.toml` will automatically record all HTTP requests made and assert
on them.

Builds on top of https://github.com/databricks/cli/pull/2226

## Tests
Modifying existing acceptance test.
---
 acceptance/acceptance_test.go                 | 37 ++++++++++++++++++-
 acceptance/config_test.go                     |  4 ++
 .../workspace/jobs/create/out.requests.txt    |  1 +
 acceptance/workspace/jobs/create/test.toml    |  2 +
 libs/testserver/server.go                     | 25 +++++++++++++
 5 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 acceptance/workspace/jobs/create/out.requests.txt

diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go
index 60f7945df..bb1d0f44f 100644
--- a/acceptance/acceptance_test.go
+++ b/acceptance/acceptance_test.go
@@ -2,6 +2,7 @@ package acceptance_test
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -219,8 +220,21 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 
 	// Start a new server with a custom configuration if the acceptance test
 	// specifies a custom server stubs.
-	if len(config.Server) > 0 {
-		server := testserver.New(t)
+	var server *testserver.Server
+
+	// Start a new server for this test if either:
+	// 1. A custom server spec is defined in the test configuration.
+	// 2. The test is configured to record requests and assert on them. We need
+	//    a duplicate of the default server to record requests because the default
+	//    server otherwise is a shared resource.
+	if len(config.Server) > 0 || config.RecordRequests {
+		server = testserver.New(t)
+		server.RecordRequests = config.RecordRequests
+
+		// If no custom server stubs are defined, add the default handlers.
+		if len(config.Server) == 0 {
+			AddHandlers(server)
+		}
 
 		for _, stub := range config.Server {
 			require.NotEmpty(t, stub.Pattern)
@@ -249,6 +263,25 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	cmd.Dir = tmpDir
 	err = cmd.Run()
 
+	// Write the requests made to the server to a output file if the test is
+	// configured to record requests.
+	if config.RecordRequests {
+		f, err := os.OpenFile(filepath.Join(tmpDir, "out.requests.txt"), os.O_CREATE|os.O_WRONLY, 0o644)
+		require.NoError(t, err)
+
+		for _, req := range server.Requests {
+			reqJson, err := json.Marshal(req)
+			require.NoError(t, err)
+
+			line := fmt.Sprintf("%s\n", reqJson)
+			_, err = f.WriteString(line)
+			require.NoError(t, err)
+		}
+
+		err = f.Close()
+		require.NoError(t, err)
+	}
+
 	// Include exit code in output (if non-zero)
 	formatOutput(out, err)
 	require.NoError(t, out.Close())
diff --git a/acceptance/config_test.go b/acceptance/config_test.go
index f340f0367..beceb6a08 100644
--- a/acceptance/config_test.go
+++ b/acceptance/config_test.go
@@ -40,6 +40,10 @@ type TestConfig struct {
 	// }
 	// '''
 	Server []ServerStub
+
+	// Record the requests made to the server and write them as output to
+	// out.requests.txt
+	RecordRequests bool
 }
 
 type ServerStub struct {
diff --git a/acceptance/workspace/jobs/create/out.requests.txt b/acceptance/workspace/jobs/create/out.requests.txt
new file mode 100644
index 000000000..b22876b70
--- /dev/null
+++ b/acceptance/workspace/jobs/create/out.requests.txt
@@ -0,0 +1 @@
+{"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}}
diff --git a/acceptance/workspace/jobs/create/test.toml b/acceptance/workspace/jobs/create/test.toml
index 94e5eee13..1bf36547b 100644
--- a/acceptance/workspace/jobs/create/test.toml
+++ b/acceptance/workspace/jobs/create/test.toml
@@ -1,3 +1,5 @@
+RecordRequests = true
+
 [[Server]]
 Pattern = "POST /api/2.1/jobs/create"
 Response.Body = '''
diff --git a/libs/testserver/server.go b/libs/testserver/server.go
index 9ebfe3ba0..2e8dbdfda 100644
--- a/libs/testserver/server.go
+++ b/libs/testserver/server.go
@@ -2,9 +2,12 @@ package testserver
 
 import (
 	"encoding/json"
+	"io"
 	"net/http"
 	"net/http/httptest"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/databricks/cli/internal/testutil"
 )
 
@@ -13,6 +16,16 @@ type Server struct {
 	Mux *http.ServeMux
 
 	t testutil.TestingT
+
+	RecordRequests bool
+
+	Requests []Request
+}
+
+type Request struct {
+	Method string `json:"method"`
+	Path   string `json:"path"`
+	Body   any    `json:"body"`
 }
 
 func New(t testutil.TestingT) *Server {
@@ -37,6 +50,18 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) {
 			return
 		}
 
+		if s.RecordRequests {
+			body, err := io.ReadAll(r.Body)
+			assert.NoError(s.t, err)
+
+			s.Requests = append(s.Requests, Request{
+				Method: r.Method,
+				Path:   r.URL.Path,
+				Body:   json.RawMessage(body),
+			})
+
+		}
+
 		w.Header().Set("Content-Type", "application/json")
 
 		var respBytes []byte