mockgcp: support VertexAIDashboard

Author: justinsb

config/tests/samples/create/harness.go

@@ -449,6 +449,8 @@ func MaybeSkip(t *testing.T, name string, resources []*unstructured.Unstructured
 
 			case schema.GroupKind{Group: "tags.cnrm.cloud.google.com", Kind: "TagsTagKey"}:
 
+			case schema.GroupKind{Group: "vertexai.cnrm.cloud.google.com", Kind: "VertexAITensorboard"}:
+
 			default:
 				t.Skipf("gk %v not suppported by mock gcp; skipping", gvk.GroupKind())
 			}

mockgcp/common/operations/http.go

@@ -27,15 +27,18 @@ import (
 	"k8s.io/klog/v2"
 )
 
-func (s *Operations) RegisterOperationsHandler(prefix string) func(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error {
+func (s *Operations) RegisterOperationsPath(path string) func(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error {
 	return func(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error {
 		forwardResponseOptions := mux.GetForwardResponseOptions()
 
-		// GET /{prefix}/operations/{name}
-		if err := mux.HandlePath("GET", "/"+prefix+"/operations/{name}", func(w http.ResponseWriter, r *http.Request, pathParams map[string]string) {
+		if err := mux.HandlePath("GET", path, func(w http.ResponseWriter, r *http.Request, pathParams map[string]string) {
 			ctx := r.Context()
 			name := pathParams["name"]
+			prefix := pathParams["prefix"]
 			req := &longrunningpb.GetOperationRequest{Name: "operations/" + name}
+			if prefix != "" {
+				req.Name = prefix + "/operations/" + name
+			}
 			op, err := s.GetOperation(ctx, req)
 			if err != nil {
 				if status.Code(err) == codes.NotFound {

mockgcp/common/operations/operations.go

@@ -62,14 +62,17 @@ func (s *Operations) NewLRO(ctx context.Context) (*pb.Operation, error) {
 	return op, nil
 }
 
-func (s *Operations) StartLRO(ctx context.Context, metadata proto.Message, callback func() (proto.Message, error)) (*pb.Operation, error) {
+func (s *Operations) StartLRO(ctx context.Context, prefix string, metadata proto.Message, callback func() (proto.Message, error)) (*pb.Operation, error) {
 	now := time.Now()
 	millis := now.UnixMilli()
 	id := uuid.NewUUID()
 
 	op := &pb.Operation{}
 
 	op.Name = fmt.Sprintf("operations/operation-%d-%s", millis, id)
+	if prefix != "" {
+		op.Name = prefix + "/" + op.Name
+	}
 	op.Done = false
 
 	if metadata != nil {
@@ -95,26 +98,10 @@ func (s *Operations) StartLRO(ctx context.Context, metadata proto.Message, callb
 			return
 		}
 
-		finished.Done = true
-		if err != nil {
-			finished.Result = &pb.Operation_Error{
-				Error: &rpcstatus.Status{
-					Message: fmt.Sprintf("error processing operation: %v", err),
-				},
-			}
-		} else {
-			resultAny, err := anypb.New(result)
-			if err != nil {
-				klog.Warningf("error building anypb for result: %v", err)
-				finished.Result = &pb.Operation_Response{}
-			} else {
-				rewriteTypes(resultAny)
-
-				finished.Result = &pb.Operation_Response{
-					Response: resultAny,
-				}
-			}
+		if err2 := markDone(finished, result, err); err2 != nil {
+			klog.Warningf("error marking LRO as done: %v", err2)
 		}
+
 		if err := s.storage.Update(ctx, fqn, finished); err != nil {
 			klog.Warningf("error updating LRO: %v", err)
 			return
@@ -124,6 +111,65 @@ func (s *Operations) StartLRO(ctx context.Context, metadata proto.Message, callb
 	return op, nil
 }
 
+func markDone(op *pb.Operation, result proto.Message, err error) error {
+	op.Done = true
+	if err != nil {
+		op.Result = &pb.Operation_Error{
+			Error: &rpcstatus.Status{
+				Message: fmt.Sprintf("error processing operation: %v", err),
+			},
+		}
+	} else {
+		resultAny, err := anypb.New(result)
+		if err != nil {
+			klog.Warningf("error building anypb for result: %v", err)
+			op.Result = &pb.Operation_Response{}
+		} else {
+			rewriteTypes(resultAny)
+
+			op.Result = &pb.Operation_Response{
+				Response: resultAny,
+			}
+		}
+	}
+	return nil
+}
+
+func (s *Operations) DoneLRO(ctx context.Context, prefix string, metadata proto.Message, result proto.Message) (*pb.Operation, error) {
+	now := time.Now()
+	millis := now.UnixMilli()
+	id := uuid.NewUUID()
+
+	op := &pb.Operation{}
+
+	op.Name = fmt.Sprintf("operations/operation-%d-%s", millis, id)
+	if prefix != "" {
+		op.Name = prefix + "/" + op.Name
+	}
+	op.Done = false
+
+	if err := markDone(op, result, nil); err != nil {
+		return nil, err
+	}
+
+	if metadata != nil {
+		metadataAny, err := anypb.New(metadata)
+		if err != nil {
+			return nil, fmt.Errorf("error building anypb for metadata: %w", err)
+		}
+		rewriteTypes(metadataAny)
+
+		op.Metadata = metadataAny
+	}
+	fqn := op.Name
+
+	if err := s.storage.Create(ctx, fqn, op); err != nil {
+		return nil, status.Errorf(codes.Internal, "error creating LRO: %v", err)
+	}
+
+	return op, nil
+}
+
 func rewriteTypes(any *anypb.Any) {
 	// Fix our mockgcp hack
 	if strings.HasPrefix(any.TypeUrl, "type.googleapis.com/mockgcp.") {

mockgcp/common/projects/projects.go

@@ -46,6 +46,8 @@ func (n *ProjectName) String() string {
 	return "projects/" + n.OriginalValue
 }
 
+// ParseProjectName parses a string into a ProjectName.
+// The expected form is projects/<projectIDOrNumber>
 func ParseProjectName(name string) (*ProjectName, error) {
 	tokens := strings.Split(name, "/")
 	if len(tokens) == 2 && tokens[0] == "projects" {

mockgcp/mock_http_roundtrip.go

@@ -32,6 +32,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/common"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/mockaiplatform"
 	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/mockapikeys"
 	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/mockbilling"
 	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/mockcertificatemanager"
@@ -88,6 +89,7 @@ func NewMockRoundTripper(t *testing.T, k8sClient client.Client, storage storage.
 	var services []MockService
 
 	services = append(services, resourcemanagerService)
+	services = append(services, mockaiplatform.New(env, storage))
 	services = append(services, mockapikeys.New(env, storage))
 	services = append(services, mockbilling.New(env, storage))
 	services = append(services, mockcertificatemanager.New(env, storage))

mockgcp/mockaiplatform/service.go

@@ -0,0 +1,69 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mockaiplatform
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/common"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/common/httpmux"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/common/operations"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/common/projects"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/pkg/storage"
+	"google.golang.org/grpc"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	pb "github.com/GoogleCloudPlatform/k8s-config-connector/mockgcp/generated/mockgcp/cloud/aiplatform/v1beta1"
+)
+
+// MockService represents a mocked aiplatform service.
+type MockService struct {
+	kube    client.Client
+	storage storage.Storage
+
+	projects   projects.ProjectStore
+	operations *operations.Operations
+}
+
+// New creates a MockService.
+func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
+	s := &MockService{
+		kube:       env.GetKubeClient(),
+		storage:    storage,
+		projects:   env.GetProjects(),
+		operations: operations.NewOperationsService(storage),
+	}
+	return s
+}
+
+func (s *MockService) ExpectedHost() string {
+	// TODO: Support more endpoints
+	return "us-central1-aiplatform.googleapis.com"
+}
+
+func (s *MockService) Register(grpcServer *grpc.Server) {
+	pb.RegisterTensorboardServiceServer(grpcServer, &tensorboardService{MockService: s})
+}
+
+func (s *MockService) NewHTTPMux(ctx context.Context, conn *grpc.ClientConn) (http.Handler, error) {
+	mux, err := httpmux.NewServeMux(ctx, conn, pb.RegisterTensorboardServiceHandler,
+		s.operations.RegisterOperationsPath("/v1beta1/{prefix=**}/operations/{name}"))
+	if err != nil {
+		return nil, err
+	}
+
+	return mux, nil
+}