Merge pull request GoogleCloudPlatform#1254 from justinsb/vertexai_integration

mockgcp: vertexai support

Author: google-oss-prow[bot]

config/tests/samples/create/harness.go

@@ -453,6 +453,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 %v; skipping", gvk.GroupKind(), name)
 			}

mockgcp/Makefile

@@ -30,6 +30,7 @@ gen-proto:
 		./third_party/googleapis/mockgcp/api/apikeys/v2/*.proto \
 		./third_party/googleapis/mockgcp/storage/v1/*.proto \
 		./third_party/googleapis/mockgcp/iam/admin/v1/*.proto \
+		./third_party/googleapis/mockgcp/cloud/aiplatform/v1beta1/*.proto \
 		./third_party/googleapis/mockgcp/cloud/billing/v1/*.proto \
 		./third_party/googleapis/mockgcp/cloud/certificatemanager/v1/*.proto \
 		./third_party/googleapis/mockgcp/cloud/compute/v1/*.proto \

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" {