github.com/databricks/cli@v0.203.0/internal/acc/workspace.go (about)

     1  package acc
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/databricks/databricks-sdk-go"
     8  	"github.com/databricks/databricks-sdk-go/service/compute"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  type WorkspaceT struct {
    13  	*testing.T
    14  
    15  	W *databricks.WorkspaceClient
    16  
    17  	ctx context.Context
    18  
    19  	exec *compute.CommandExecutorV2
    20  }
    21  
    22  func WorkspaceTest(t *testing.T) (context.Context, *WorkspaceT) {
    23  	loadDebugEnvIfRunFromIDE(t, "workspace")
    24  
    25  	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
    26  
    27  	w, err := databricks.NewWorkspaceClient()
    28  	require.NoError(t, err)
    29  
    30  	wt := &WorkspaceT{
    31  		T: t,
    32  
    33  		W: w,
    34  
    35  		ctx: context.Background(),
    36  	}
    37  
    38  	return wt.ctx, wt
    39  }
    40  
    41  func (t *WorkspaceT) TestClusterID() string {
    42  	clusterID := GetEnvOrSkipTest(t.T, "TEST_BRICKS_CLUSTER_ID")
    43  	err := t.W.Clusters.EnsureClusterIsRunning(t.ctx, clusterID)
    44  	require.NoError(t, err)
    45  	return clusterID
    46  }
    47  
    48  func (t *WorkspaceT) RunPython(code string) (string, error) {
    49  	var err error
    50  
    51  	// Create command executor only once per test.
    52  	if t.exec == nil {
    53  		t.exec, err = t.W.CommandExecution.Start(t.ctx, t.TestClusterID(), compute.LanguagePython)
    54  		require.NoError(t, err)
    55  
    56  		t.Cleanup(func() {
    57  			err := t.exec.Destroy(t.ctx)
    58  			require.NoError(t, err)
    59  		})
    60  	}
    61  
    62  	results, err := t.exec.Execute(t.ctx, code)
    63  	require.NoError(t, err)
    64  	require.NotEqual(t, compute.ResultTypeError, results.ResultType, results.Cause)
    65  	output, ok := results.Data.(string)
    66  	require.True(t, ok, "unexpected type %T", results.Data)
    67  	return output, nil
    68  }