Reduce integration test cluster capacity usage

NEXT_CHANGELOG.md

@@ -12,4 +12,6 @@
 
 ### Internal Changes
 
+- Reduce integration test cluster usage by replacing waiter coverage with HTTP fixtures and removing redundant live tests.
+
 ### API Changes

internal/clusters_test.go

@@ -2,60 +2,94 @@ package internal
 
 import (
 	"context"
-	"strings"
+	"net/http"
 	"testing"
 	"time"
 
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/databricks/databricks-sdk-go/common/environment"
+	"github.com/databricks/databricks-sdk-go/qa"
 	"github.com/databricks/databricks-sdk-go/retries"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func sharedRunningCluster(t *testing.T, ctx context.Context,
-	w *databricks.WorkspaceClient,
-) string {
-	clusterId := GetEnvOrSkipTest(t, "TEST_GO_SDK_CLUSTER_ID")
-	err := w.Clusters.EnsureClusterIsRunning(ctx, clusterId)
-	require.NoError(t, err)
-	return clusterId
-}
+func TestClustersCreateTimesOutAndCleansUp(t *testing.T) {
+	const clusterID = "test-cluster"
+	createRequest := compute.CreateCluster{
+		ClusterName:            "test-cluster",
+		SparkVersion:           "15.4.x-scala2.12",
+		InstancePoolId:         "test-pool",
+		AutoterminationMinutes: 10,
+		NumWorkers:             1,
+	}
+	pending := compute.ClusterDetails{
+		ClusterId:    clusterID,
+		State:        compute.StatePending,
+		StateMessage: "still provisioning",
+	}
 
-func TestAccClustersCreateFailsWithTimeoutNoTranspile(t *testing.T) {
-	ctx, w := workspaceTest(t)
+	qa.HTTPFixtures{
+		{
+			Method:          http.MethodPost,
+			Resource:        "/api/2.1/clusters/create",
+			ExpectedRequest: createRequest,
+			Response: compute.CreateClusterResponse{
+				ClusterId: clusterID,
+			},
+		},
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.1/clusters/get?cluster_id=" + clusterID,
+			Response: pending,
+		},
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.1/clusters/get?cluster_id=" + clusterID,
+			Response: pending,
+		},
+		{
+			Method:   http.MethodPost,
+			Resource: "/api/2.1/clusters/delete",
+			ExpectedRequest: compute.DeleteCluster{
+				ClusterId: clusterID,
+			},
+		},
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.1/clusters/get?cluster_id=" + clusterID,
+			Response: compute.ClusterDetails{
+				ClusterId: clusterID,
+				State:     compute.StateTerminated,
+			},
+		},
+	}.ApplyClient(t, func(ctx context.Context, apiClient *client.DatabricksClient) {
+		clusters := compute.NewClusters(apiClient)
+		createCtx, cancelCreate := context.WithCancel(ctx)
+		defer cancelCreate()
 
-	// Fetch list of spark runtime versions
-	sparkVersions, err := w.Clusters.SparkVersions(ctx)
-	require.NoError(t, err)
+		var gotStates []compute.State
+		_, err := clusters.CreateAndWait(
+			createCtx,
+			createRequest,
+			retries.Timeout[compute.ClusterDetails](time.Minute),
+			retries.OnPoll(func(details *compute.ClusterDetails) {
+				gotStates = append(gotStates, details.State)
+				if len(gotStates) == 2 {
+					cancelCreate()
+				}
+			}),
+		)
+		assert.EqualError(t, err, "timed out: still provisioning")
+		assert.Equal(t, []compute.State{compute.StatePending, compute.StatePending}, gotStates)
 
-	// Select the latest LTS version without Photon
-	latest, err := sparkVersions.Select(compute.SparkVersionRequest{
-		Latest:          true,
-		LongTermSupport: true,
+		deleted, err := clusters.DeleteByClusterIdAndWait(ctx, clusterID)
+		require.NoError(t, err)
+		assert.Equal(t, clusterID, deleted.ClusterId)
+		assert.Equal(t, compute.StateTerminated, deleted.State)
 	})
-	require.NoError(t, err)
-
-	var clusterId string
-
-	// Create a cluster with unreasonably low timeout
-	_, err = w.Clusters.CreateAndWait(ctx, compute.CreateCluster{
-		ClusterName:            RandomName(t.Name()),
-		SparkVersion:           latest,
-		InstancePoolId:         GetEnvOrSkipTest(t, "TEST_INSTANCE_POOL_ID"),
-		AutoterminationMinutes: 10,
-		NumWorkers:             1,
-	}, retries.Timeout[compute.ClusterDetails](15*time.Second),
-		func(i *retries.Info[compute.ClusterDetails]) {
-			if i.Info == nil {
-				return
-			}
-			clusterId = i.Info.ClusterId
-		})
-	assert.True(t, strings.HasPrefix(err.Error(), "timed out: "))
-	_, err = w.Clusters.DeleteByClusterIdAndWait(ctx, clusterId)
-	require.NoError(t, err)
 }
 
 func TestAccClustersGetCorrectErrorMessageNoTranspile(t *testing.T) {

internal/commands_test.go

@@ -30,44 +30,3 @@ func TestAccCommands(t *testing.T) {
 	require.NoError(t, err)
 	assert.EqualError(t, res.Err(), "ZeroDivisionError: division by zero")
 }
-
-func TestAccCommandsDirectUsage(t *testing.T) {
-	ctx, w := workspaceTest(t)
-
-	clusterId := GetEnvOrSkipTest(t, "TEST_DEFAULT_CLUSTER_ID")
-	err := w.Clusters.EnsureClusterIsRunning(ctx, clusterId)
-	require.NoError(t, err)
-
-	context, err := w.CommandExecution.CreateAndWait(ctx, compute.CreateContext{
-		ClusterId: clusterId,
-		Language:  compute.LanguagePython,
-	})
-	require.NoError(t, err)
-
-	t.Cleanup(func() {
-		err = w.CommandExecution.Destroy(ctx, compute.DestroyContext{
-			ClusterId: clusterId,
-			ContextId: context.Id,
-		})
-		require.NoError(t, err)
-	})
-
-	textResults, err := w.CommandExecution.ExecuteAndWait(ctx, compute.Command{
-		ClusterId: clusterId,
-		ContextId: context.Id,
-		Language:  compute.LanguagePython,
-		Command:   "print(1)",
-	})
-	require.NoError(t, err)
-
-	assert.Equal(t, "1", textResults.Results.Text())
-
-	failingCommand, err := w.CommandExecution.ExecuteAndWait(ctx, compute.Command{
-		ClusterId: clusterId,
-		ContextId: context.Id,
-		Language:  compute.LanguagePython,
-		Command:   "1/0",
-	})
-	require.NoError(t, err)
-	assert.EqualError(t, failingCommand.Results.Err(), "ZeroDivisionError: division by zero")
-}

internal/libraries_test.go

@@ -1,37 +1,98 @@
 package internal
 
 import (
+	"context"
+	"net/http"
 	"testing"
+	"time"
 
+	"github.com/databricks/databricks-sdk-go/client"
+	"github.com/databricks/databricks-sdk-go/qa"
+	"github.com/databricks/databricks-sdk-go/retries"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/stretchr/testify/require"
 )
 
-func TestAccLibraries(t *testing.T) {
-	ctx, w := workspaceTest(t)
-	clusterId := sharedRunningCluster(t, ctx, w)
+func TestLibrariesUpdateAndWait(t *testing.T) {
+	const clusterID = "test-cluster"
+	library := compute.Library{
+		Pypi: &compute.PythonPyPiLibrary{
+			Package: "test-package",
+		},
+	}
+	installRequest := compute.InstallLibraries{
+		ClusterId: clusterID,
+		Libraries: []compute.Library{library},
+	}
+	uninstallRequest := compute.UninstallLibraries{
+		ClusterId: clusterID,
+		Libraries: []compute.Library{library},
+	}
 
-	err := w.Libraries.UpdateAndWait(ctx, compute.Update{
-		ClusterId: clusterId,
-		Install: []compute.Library{
-			{
-				Pypi: &compute.PythonPyPiLibrary{
-					Package: "dbl-tempo",
+	qa.HTTPFixtures{
+		{
+			Method:          http.MethodPost,
+			Resource:        "/api/2.0/libraries/install",
+			ExpectedRequest: installRequest,
+		},
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.0/libraries/cluster-status?cluster_id=" + clusterID,
+			Response: compute.ClusterLibraryStatuses{
+				ClusterId: clusterID,
+				LibraryStatuses: []compute.LibraryFullStatus{
+					{
+						Library: &library,
+						Status:  compute.LibraryInstallStatusPending,
+					},
 				},
 			},
 		},
-	})
-	require.NoError(t, err)
-
-	err = w.Libraries.UpdateAndWait(ctx, compute.Update{
-		ClusterId: clusterId,
-		Uninstall: []compute.Library{
-			{
-				Pypi: &compute.PythonPyPiLibrary{
-					Package: "dbl-tempo",
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.0/libraries/cluster-status?cluster_id=" + clusterID,
+			Response: compute.ClusterLibraryStatuses{
+				ClusterId: clusterID,
+				LibraryStatuses: []compute.LibraryFullStatus{
+					{
+						Library: &library,
+						Status:  compute.LibraryInstallStatusInstalled,
+					},
+				},
+			},
+		},
+		{
+			Method:          http.MethodPost,
+			Resource:        "/api/2.0/libraries/uninstall",
+			ExpectedRequest: uninstallRequest,
+		},
+		{
+			Method:   http.MethodGet,
+			Resource: "/api/2.0/libraries/cluster-status?cluster_id=" + clusterID,
+			Response: compute.ClusterLibraryStatuses{
+				ClusterId: clusterID,
+				LibraryStatuses: []compute.LibraryFullStatus{
+					{
+						Library: &library,
+						Status:  compute.LibraryInstallStatusUninstallOnRestart,
+					},
 				},
 			},
 		},
+	}.ApplyClient(t, func(ctx context.Context, apiClient *client.DatabricksClient) {
+		libraries := compute.NewLibraries(apiClient)
+		timeout := retries.Timeout[compute.ClusterLibraryStatuses](5 * time.Second)
+
+		err := libraries.UpdateAndWait(ctx, compute.Update{
+			ClusterId: clusterID,
+			Install:   []compute.Library{library},
+		}, timeout)
+		require.NoError(t, err)
+
+		err = libraries.UpdateAndWait(ctx, compute.Update{
+			ClusterId: clusterID,
+			Uninstall: []compute.Library{library},
+		}, timeout)
+		require.NoError(t, err)
 	})
-	require.NoError(t, err)
 }