mirror of https://github.com/databricks/cli.git
130 lines
2.8 KiB
Go
130 lines
2.8 KiB
Go
package acceptance_test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/databricks/databricks-sdk-go/service/compute"
|
|
"github.com/databricks/databricks-sdk-go/service/iam"
|
|
"github.com/databricks/databricks-sdk-go/service/workspace"
|
|
)
|
|
|
|
type TestServer struct {
|
|
*httptest.Server
|
|
Mux *http.ServeMux
|
|
Port int
|
|
}
|
|
|
|
type HandlerFunc func(r *http.Request) (any, error)
|
|
|
|
func NewTestServer() *TestServer {
|
|
mux := http.NewServeMux()
|
|
server := httptest.NewServer(mux)
|
|
port := server.Listener.Addr().(*net.TCPAddr).Port
|
|
|
|
return &TestServer{
|
|
Server: server,
|
|
Mux: mux,
|
|
Port: port,
|
|
}
|
|
}
|
|
|
|
func (s *TestServer) Handle(pattern string, handler HandlerFunc) {
|
|
s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
|
|
resp, err := handler(r)
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
|
|
var respBytes []byte
|
|
|
|
respString, ok := resp.(string)
|
|
if ok {
|
|
respBytes = []byte(respString)
|
|
} else {
|
|
respBytes, err = json.MarshalIndent(resp, "", " ")
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
}
|
|
|
|
if _, err := w.Write(respBytes); err != nil {
|
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
})
|
|
}
|
|
|
|
func StartServer(t *testing.T) *TestServer {
|
|
server := NewTestServer()
|
|
t.Cleanup(func() {
|
|
server.Close()
|
|
})
|
|
return server
|
|
}
|
|
|
|
func AddHandlers(server *TestServer) {
|
|
server.Handle("/api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
|
|
return compute.ListPoliciesResponse{
|
|
Policies: []compute.Policy{
|
|
{
|
|
PolicyId: "5678",
|
|
Name: "wrong-cluster-policy",
|
|
},
|
|
{
|
|
PolicyId: "9876",
|
|
Name: "some-test-cluster-policy",
|
|
},
|
|
},
|
|
}, nil
|
|
})
|
|
|
|
server.Handle("/api/2.0/instance-pools/list", func(r *http.Request) (any, error) {
|
|
return compute.ListInstancePools{
|
|
InstancePools: []compute.InstancePoolAndStats{
|
|
{
|
|
InstancePoolName: "some-test-instance-pool",
|
|
InstancePoolId: "1234",
|
|
},
|
|
},
|
|
}, nil
|
|
})
|
|
|
|
server.Handle("/api/2.1/clusters/list", func(r *http.Request) (any, error) {
|
|
return compute.ListClustersResponse{
|
|
Clusters: []compute.ClusterDetails{
|
|
{
|
|
ClusterName: "some-test-cluster",
|
|
ClusterId: "4321",
|
|
},
|
|
{
|
|
ClusterName: "some-other-cluster",
|
|
ClusterId: "9876",
|
|
},
|
|
},
|
|
}, nil
|
|
})
|
|
|
|
server.Handle("/api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, error) {
|
|
return iam.User{
|
|
UserName: "tester@databricks.com",
|
|
}, nil
|
|
})
|
|
|
|
server.Handle("/api/2.0/workspace/get-status", func(r *http.Request) (any, error) {
|
|
return workspace.ObjectInfo{
|
|
ObjectId: 1001,
|
|
ObjectType: "DIRECTORY",
|
|
Path: "",
|
|
ResourceId: "1001",
|
|
}, nil
|
|
})
|
|
}
|