From 42c06267eb64b2ccbe55985f4a426335e404cd3b Mon Sep 17 00:00:00 2001
From: Serge Smertin <259697+nfx@users.noreply.github.com>
Date: Mon, 11 Dec 2023 08:30:19 -0800
Subject: [PATCH] Stub out Python virtual environment installation for `labs`
 commands (#1057)

This PR removes 15 seconds from `make test` runtime
---
 cmd/labs/project/installer_test.go | 10 +++++++++-
 libs/process/stub.go               | 22 +++++++++++++++++++---
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/cmd/labs/project/installer_test.go b/cmd/labs/project/installer_test.go
index 60af43c6..709e14f2 100644
--- a/cmd/labs/project/installer_test.go
+++ b/cmd/labs/project/installer_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/databricks/cli/cmd/labs/project"
 	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/cli/libs/python"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -194,6 +195,13 @@ func TestInstallerWorksForReleases(t *testing.T) {
 
 	ctx := installerContext(t, server)
 
+	ctx, stub := process.WithStub(ctx)
+	stub.WithStdoutFor(`python[\S]+ --version`, "Python 3.10.5")
+	// on Unix, we call `python3`, but on Windows it is `python.exe`
+	stub.WithStderrFor(`python[\S]+ -m venv .*/.databricks/labs/blueprint/state/venv`, "[mock venv create]")
+	stub.WithStderrFor(`python[\S]+ -m pip install .`, "[mock pip install]")
+	stub.WithStdoutFor(`python[\S]+ install.py`, "setting up important infrastructure")
+
 	// simulate the case of GitHub Actions
 	ctx = env.Set(ctx, "DATABRICKS_HOST", server.URL)
 	ctx = env.Set(ctx, "DATABRICKS_TOKEN", "...")
@@ -228,7 +236,7 @@ func TestInstallerWorksForReleases(t *testing.T) {
 	//     │               │   │       └── site-packages
 	//     │               │   │           ├── ...
 	//     │               │   │           ├── distutils-precedence.pth
-	r := internal.NewCobraTestRunnerWithContext(t, ctx, "labs", "install", "blueprint")
+	r := internal.NewCobraTestRunnerWithContext(t, ctx, "labs", "install", "blueprint", "--debug")
 	r.RunAndExpectOutput("setting up important infrastructure")
 }
 
diff --git a/libs/process/stub.go b/libs/process/stub.go
index 280a9a8a..8472f65d 100644
--- a/libs/process/stub.go
+++ b/libs/process/stub.go
@@ -2,9 +2,11 @@ package process
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"strings"
 )
 
@@ -128,14 +130,24 @@ func (s *processStub) normCmd(v *exec.Cmd) string {
 	// "/var/folders/bc/7qf8yghj6v14t40096pdcqy40000gp/T/tmp.03CAcYcbOI/python3" becomes "python3".
 	// Use [processStub.WithCallback] if you need to match against the full executable path.
 	binaryName := filepath.Base(v.Path)
-	args := strings.Join(v.Args[1:], " ")
+	var unixArgs []string
+	for _, arg := range v.Args[1:] {
+		unixArgs = append(unixArgs, filepath.ToSlash(arg))
+	}
+	args := strings.Join(unixArgs, " ")
 	return fmt.Sprintf("%s %s", binaryName, args)
 }
 
+var ErrStubContinue = errors.New("continue executing the stub after callback")
+
 func (s *processStub) run(cmd *exec.Cmd) error {
 	s.calls = append(s.calls, cmd)
-	resp, ok := s.responses[s.normCmd(cmd)]
-	if ok {
+	for pattern, resp := range s.responses {
+		re := regexp.MustCompile(pattern)
+		norm := s.normCmd(cmd)
+		if !re.MatchString(norm) {
+			continue
+		}
 		if resp.stdout != "" {
 			cmd.Stdout.Write([]byte(resp.stdout))
 		}
@@ -147,6 +159,10 @@ func (s *processStub) run(cmd *exec.Cmd) error {
 	if s.callback != nil {
 		return s.callback(cmd)
 	}
+	var zeroStub reponseStub
+	if s.reponseStub == zeroStub {
+		return fmt.Errorf("no default process stub")
+	}
 	if s.reponseStub.stdout != "" {
 		cmd.Stdout.Write([]byte(s.reponseStub.stdout))
 	}