Support locking in integration tests

go.mod

@@ -4,6 +4,7 @@ go 1.18
 
 require (
 	github.com/google/go-querystring v1.1.0
+	github.com/google/uuid v1.4.0
 	github.com/stretchr/testify v1.8.4
 	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e
 	golang.org/x/mod v0.14.0

go.sum

@@ -54,6 +54,7 @@ github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
 github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
+github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfFxPRy3Bf7vr3h0cechB90XaQs=
 github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
 github.com/googleapis/gax-go/v2 v2.12.0 h1:A+gCJKdRfqXkr+BIRGtZLibNXf0m1f9E4HG56etFpas=

internal/git_credentials_test.go

@@ -1,16 +1,32 @@
 package internal
 
 import (
+	"context"
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/qa/lock"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func acquireGitCredentialsLock(ctx context.Context, t *testing.T, w *databricks.WorkspaceClient) {
+	me, err := w.CurrentUser.Me(ctx)
+	require.NoError(t, err)
+	lockable := lock.GitCredentials{
+		WorkspaceHost: w.Config.Host,
+		Username:      me.UserName,
+	}
+	_, err = lock.Acquire(ctx, lockable, lock.InTest(t))
+	require.NoError(t, err)
+}
+
 func TestAccGitCredentials(t *testing.T) {
 	ctx, w := workspaceTest(t)
 
+	// skip-next-line-roll
+	acquireGitCredentialsLock(ctx, t, w)
 	list, err := w.GitCredentials.ListAll(ctx)
 	require.NoError(t, err)
 	for _, v := range list {

internal/init_test.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"math/rand"
 	"os"
-	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -16,6 +15,7 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/logger"
+	"github.com/databricks/databricks-sdk-go/qa"
 )
 
 const fullCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
@@ -147,23 +147,17 @@ func RandomHex(prefix string, randLen int) string {
 }
 
 func skipf(t *testing.T) func(format string, args ...any) {
-	if isInDebug() {
+	if qa.IsInDebug() {
 		// VSCode "debug test" feature doesn't show dlv logs,
 		// so that we fail here for maintainer productivity.
 		return t.Fatalf
 	}
 	return t.Skipf
 }
 
-// detects if test is run from "debug test" feature in VSCode
-func isInDebug() bool {
-	ex, _ := os.Executable()
-	return strings.HasPrefix(path.Base(ex), "__debug_bin")
-}
-
 // loads debug environment from ~/.databricks/debug-env.json
 func loadDebugEnvIfRunsFromIDE(t *testing.T, key string) {
-	if !isInDebug() {
+	if !qa.IsInDebug() {
 		return
 	}
 	home, err := os.UserHomeDir()

openapi/roll/tool.go

@@ -353,6 +353,47 @@ func (s *Suite) expectExamples(file *ast.File) {
 	}
 }
 
+// Given a function declaration, return a map of statement index to comment
+// immediately preceding the statement.
+//
+// Example:
+//
+// func foo() {
+//   // comment 1
+//   bar()
+//   baz()
+//   // comment 2
+//   baz()
+// }
+//
+// Returns:
+//
+// map[int]string{
+//   0: "comment 1",
+//   2: "comment 2",
+// }
+func (s *Suite) getCommentPrecedingStatementMap(fn *ast.FuncDecl, file *ast.File) map[int]string {
+	res := map[int]string{}
+	commentIndex := 0
+	for _, cg := range file.Comments {
+		if cg.End() >= fn.Pos() {
+			break
+		}
+		commentIndex += 1
+	}
+	for i, stmt := range fn.Body.List {
+		if commentIndex >= len(file.Comments) {
+			return res
+		}
+		commentGroup := file.Comments[commentIndex]
+		if stmt.Pos() > commentGroup.End() {
+			res[i] = strings.TrimSpace(commentGroup.Text())
+			commentIndex += 1
+		}
+	}
+	return res
+}
+
 func (s *Suite) expectFn(fn *ast.FuncDecl, file *ast.File) *example {
 	testName := fn.Name.Name
 	testName = strings.TrimPrefix(testName, "TestAcc")
@@ -364,19 +405,23 @@ func (s *Suite) expectFn(fn *ast.FuncDecl, file *ast.File) *example {
 		},
 		scope: map[string]expression{},
 	}
-	lastPos := fn.Pos()
-	hint := ""
-	for _, v := range fn.Body.List {
-		for _, cmt := range file.Comments {
-			if cmt.End() < lastPos {
-				continue
-			}
-			if cmt.Pos() > v.Pos() {
+	commentMap := s.getCommentPrecedingStatementMap(fn, file)
+	for i, v := range fn.Body.List {
+		hint := commentMap[i]
+		if hint == "skip-next-line-roll" {
+			switch node := v.(type) {
+			case *ast.ExprStmt:
+				if _, ok := node.X.(*ast.CallExpr); !ok {
+					buf := &strings.Builder{}
+					ast.Fprint(buf, s.fset, node, nil)
+					panic(fmt.Errorf("skip-next-line-roll can only be used immediately before plain function call with no LHS, found %s", buf.String()))
+				}
 				continue
+			default:
+				buf := &strings.Builder{}
+				ast.Fprint(buf, s.fset, node, nil)
+				panic(fmt.Errorf("skip-next-line-roll can only be used immediately before plain function call with no LHS, found %s", buf.String()))
 			}
-			// figure out comment hint exactly above the given statement
-			hint = strings.TrimSpace(cmt.Text())
-			break
 		}
 		switch stmt := v.(type) {
 		case *ast.AssignStmt:
@@ -386,10 +431,6 @@ func (s *Suite) expectFn(fn *ast.FuncDecl, file *ast.File) *example {
 		case *ast.ExprStmt:
 			s.expectExprStmt(ex, stmt)
 		}
-		// store the end of the last statement to figure out the next
-		// statement comment.
-		lastPos = v.End()
-		hint = ""
 	}
 	return ex
 }

qa/debug.go

@@ -0,0 +1,13 @@
+package qa
+
+import (
+	"os"
+	"path"
+	"strings"
+)
+
+// detects if test is run from "debug test" feature in VSCode
+func IsInDebug() bool {
+	ex, _ := os.Executable()
+	return strings.HasPrefix(path.Base(ex), "__debug_bin")
+}

qa/lock/core/backend.go

@@ -0,0 +1,26 @@
+package core
+
+import (
+	"context"
+	"time"
+)
+
+// LockBackend is the interface for a lock backend.
+type LockBackend interface {
+	// PrepareBackend prepares the backend for use. This is called once per lock.
+	// Locking is aborted if this fails.
+	PrepareBackend(ctx context.Context, lockId string) error
+
+	// AcquireLock acquires the lock. This is called once per lock. Locking is aborted
+	// if this fails.
+	AcquireLock(ctx context.Context, contents *LockState) error
+
+	// RenewLock renews the lock. This is called periodically while the lock is held.
+	RenewLock(ctx context.Context, leaseId string) error
+
+	// ReleaseLock releases the lock. This is called once per lock.
+	ReleaseLock(ctx context.Context, leaseId string) error
+
+	// RefreshDuration returns the duration between calls to RenewLock.
+	RefreshDuration() time.Duration
+}

qa/lock/core/lockable.go

@@ -0,0 +1,7 @@
+package core
+
+// Lockable is the interface for a lockable resource.
+type Lockable interface {
+	// GetLockId returns the lock id for the resource.
+	GetLockId() string
+}

qa/lock/core/state.go

@@ -0,0 +1,101 @@
+package core
+
+import (
+	"os"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/databricks/databricks-sdk-go/qa"
+)
+
+// LockState represents the state of a lock.
+type LockState struct {
+	// Lockable is the lock id of the lockable resource.
+	Lockable string
+
+	// LeaseId is the lease id of the lock.
+	LeaseId string
+
+	// Start is the time the lock was acquired.
+	Start time.Time
+
+	// Expiry is the time the lock expires.
+	Expiry time.Time
+
+	// IsInDebug is true if the test is running in debug mode.
+	IsInDebug bool
+
+	// Test is the name of the test, if applicable.
+	Test string
+
+	// User is the user running the test.
+	User string `env:"USER"`
+
+	// GithubActions is true if the test is running in Github Actions.
+	GithubActions string `env:"GITHUB_ACTIONS"`
+
+	// GithubRef is the ref of the Github Actions workflow.
+	GithubRef string `env:"GITHUB_REF"`
+
+	// GithubRepository is the repository of the Github Actions workflow, e.g. `databricks/databricks-sdk-go`.
+	GithubRepository string `env:"GITHUB_REPOSITORY"`
+
+	// GithubRunId is the ID of the Github Actions run.
+	GithubRunId string `env:"GITHUB_RUN_ID"`
+
+	// GithubRunNumber is the number of the Github Actions run.
+	GithubRunNumber string `env:"GITHUB_RUN_NUMBER"`
+
+	// GithubSha is the SHA of the Github Actions commit.
+	GithubSha string `env:"GITHUB_SHA"`
+
+	// GithubWorkflow is the name of the Github Actions workflow.
+	GithubWorkflow string `env:"GITHUB_WORKFLOW"`
+}
+
+func (l *LockState) loadFromEnv() {
+	rv := reflect.ValueOf(l).Elem()
+	t := rv.Type()
+	for i := 0; i < rv.NumField(); i++ {
+		field := rv.Field(i)
+		tag := t.Field(i).Tag.Get("env")
+		if tag == "" {
+			continue
+		}
+		val := os.Getenv(tag)
+		if val != "" {
+			field.SetString(val)
+		}
+	}
+}
+
+// IsValid() returns true if the lock has not yet expired.
+func (l *LockState) IsValid() bool {
+	return time.Now().Before(l.Expiry)
+}
+
+// Extend() resets the lock's expiry to the current time plus the given duration.
+// This method only updates the state object; it does not actually persist the
+// change to the lock backend.
+func (l *LockState) Extend(duration time.Duration) {
+	l.Expiry = time.Now().Add(duration)
+}
+
+// NewLockState creates a new LockState object with the given lockable, lease ID,
+// and lease duration. If t is not nil, the test name is also set.
+func NewLockState(lockable Lockable, leaseId string, leaseDuration time.Duration, t *testing.T) *LockState {
+	now := time.Now()
+	state := &LockState{
+		Start:     now,
+		Expiry:    now.Add(leaseDuration),
+		Lockable:  lockable.GetLockId(),
+		LeaseId:   leaseId,
+		IsInDebug: qa.IsInDebug(),
+	}
+	state.loadFromEnv()
+	if t != nil {
+		state.Test = t.Name()
+	}
+	return state
+}