From 3c6eacb05b6980c59324b1da4ac091f22cb5733d Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Thu, 30 Jan 2025 16:13:07 +0530
Subject: [PATCH] Add feature to mock server APIs in acceptance tests (#2226)

## Changes
This PR allows us to define custom server stubs in a `test.toml` file.

Note: A followup PR will add functionality to do assertions on the API
request itself.

## Tests
New acceptance test.
---
 acceptance/acceptance_test.go               | 25 +++++++++++++++++----
 acceptance/config_test.go                   | 23 +++++++++++++++++++
 acceptance/workspace/jobs/create/output.txt |  5 +++++
 acceptance/workspace/jobs/create/script     |  1 +
 acceptance/workspace/jobs/create/test.toml  |  7 ++++++
 5 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 acceptance/workspace/jobs/create/output.txt
 create mode 100644 acceptance/workspace/jobs/create/script
 create mode 100644 acceptance/workspace/jobs/create/test.toml

diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go
index e7104a1c1..60f7945df 100644
--- a/acceptance/acceptance_test.go
+++ b/acceptance/acceptance_test.go
@@ -6,6 +6,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -112,10 +113,10 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	cloudEnv := os.Getenv("CLOUD_ENV")
 
 	if cloudEnv == "" {
-		server := testserver.New(t)
-		AddHandlers(server)
+		defaultServer := testserver.New(t)
+		AddHandlers(defaultServer)
 		// Redirect API access to local server:
-		t.Setenv("DATABRICKS_HOST", server.URL)
+		t.Setenv("DATABRICKS_HOST", defaultServer.URL)
 		t.Setenv("DATABRICKS_TOKEN", "dapi1234")
 
 		homeDir := t.TempDir()
@@ -214,6 +215,22 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 
 	args := []string{"bash", "-euo", "pipefail", EntryPointScript}
 	cmd := exec.Command(args[0], args[1:]...)
+	cmd.Env = os.Environ()
+
+	// 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)
+
+		for _, stub := range config.Server {
+			require.NotEmpty(t, stub.Pattern)
+			server.Handle(stub.Pattern, func(req *http.Request) (resp any, err error) {
+				return stub.Response.Body, nil
+			})
+		}
+		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+server.URL)
+	}
+
 	if coverDir != "" {
 		// Creating individual coverage directory for each test, because writing to the same one
 		// results in sporadic failures like this one (only if tests are running in parallel):
@@ -221,7 +238,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 		coverDir = filepath.Join(coverDir, strings.ReplaceAll(dir, string(os.PathSeparator), "--"))
 		err := os.MkdirAll(coverDir, os.ModePerm)
 		require.NoError(t, err)
-		cmd.Env = append(os.Environ(), "GOCOVERDIR="+coverDir)
+		cmd.Env = append(cmd.Env, "GOCOVERDIR="+coverDir)
 	}
 
 	// Write combined output to a file
diff --git a/acceptance/config_test.go b/acceptance/config_test.go
index 41866c4a7..f340f0367 100644
--- a/acceptance/config_test.go
+++ b/acceptance/config_test.go
@@ -29,6 +29,29 @@ type TestConfig struct {
 	// List of additional replacements to apply on this test.
 	// Old is a regexp, New is a replacement expression.
 	Repls []testdiff.Replacement
+
+	// List of server stubs to load. Example configuration:
+	//
+	// [[Server]]
+	// Pattern = "POST /api/2.1/jobs/create"
+	// Response.Body = '''
+	// {
+	// 	"job_id": 1111
+	// }
+	// '''
+	Server []ServerStub
+}
+
+type ServerStub struct {
+	// The HTTP method and path to match. Examples:
+	// 1. /api/2.0/clusters/list (matches all methods)
+	// 2. GET /api/2.0/clusters/list
+	Pattern string
+
+	// The response body to return.
+	Response struct {
+		Body string
+	}
 }
 
 // FindConfig finds the closest config file.
diff --git a/acceptance/workspace/jobs/create/output.txt b/acceptance/workspace/jobs/create/output.txt
new file mode 100644
index 000000000..a9487fe5b
--- /dev/null
+++ b/acceptance/workspace/jobs/create/output.txt
@@ -0,0 +1,5 @@
+
+>>> $CLI jobs create --json {"name":"abc"}
+{
+  "job_id":1111
+}
diff --git a/acceptance/workspace/jobs/create/script b/acceptance/workspace/jobs/create/script
new file mode 100644
index 000000000..9ff7b5b87
--- /dev/null
+++ b/acceptance/workspace/jobs/create/script
@@ -0,0 +1 @@
+trace $CLI jobs create --json '{"name":"abc"}'
diff --git a/acceptance/workspace/jobs/create/test.toml b/acceptance/workspace/jobs/create/test.toml
new file mode 100644
index 000000000..94e5eee13
--- /dev/null
+++ b/acceptance/workspace/jobs/create/test.toml
@@ -0,0 +1,7 @@
+[[Server]]
+Pattern = "POST /api/2.1/jobs/create"
+Response.Body = '''
+{
+    "job_id": 1111
+}
+'''