sigs.k8s.io/prow@v0.0.0-20240503223140-c5e374dc7eb1/test/integration/internal/fakepubsub/fakepubsub.go (about)

     1  /*
     2  Copyright 2022 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package fakepubsub
    18  
    19  import (
    20  	"context"
    21  	"encoding/json"
    22  	"fmt"
    23  	"time"
    24  
    25  	"cloud.google.com/go/pubsub"
    26  	"github.com/sirupsen/logrus"
    27  	"google.golang.org/api/option"
    28  	"google.golang.org/grpc"
    29  	"google.golang.org/grpc/credentials/insecure"
    30  
    31  	"sigs.k8s.io/prow/pkg/pubsub/subscriber"
    32  )
    33  
    34  type PubSubMessageForSub struct {
    35  	Attributes map[string]string
    36  	Data       subscriber.ProwJobEvent
    37  }
    38  
    39  type Client struct {
    40  	projectID    string
    41  	pubsubClient *pubsub.Client
    42  }
    43  
    44  func NewClient(projectID, pubsubEmulatorHost string) (*Client, error) {
    45  	client, err := newClientForEmulator(projectID, pubsubEmulatorHost)
    46  	if err != nil {
    47  		return nil, fmt.Errorf("Unable to create pubsub client to project %q for the emulator: %v", projectID, err)
    48  	}
    49  
    50  	return &Client{
    51  		projectID:    projectID,
    52  		pubsubClient: client,
    53  	}, nil
    54  }
    55  
    56  // newClientForEmulator returns a pubsub client that is hardcoded to always talk
    57  // to the fakepubsub service running in the test KIND cluster via the
    58  // pubsubEmulatorHost parameter. This is taken from
    59  // https://github.com/googleapis/google-cloud-go/blob/e43c095c94e44a95c618861f9da8f2469b53be16/pubsub/pubsub.go#L126.
    60  // This is better than getting the PUBSUB_EMULATOR_HOST environment variable
    61  // because this makes the code thread-safe (we no longer rely on a global
    62  // environment variable).
    63  func newClientForEmulator(projectID, pubsubEmulatorHost string) (*pubsub.Client, error) {
    64  	conn, err := grpc.Dial(pubsubEmulatorHost, grpc.WithTransportCredentials(insecure.NewCredentials()))
    65  	if err != nil {
    66  		return nil, fmt.Errorf("grpc.Dial: %v", err)
    67  	}
    68  	o := []option.ClientOption{option.WithGRPCConn(conn)}
    69  	o = append(o, option.WithTelemetryDisabled())
    70  	return pubsub.NewClientWithConfig(context.Background(), projectID, nil, o...)
    71  }
    72  
    73  // PublishMessage creates a Pub/Sub message that sub understands (to create a
    74  // ProwJob). The podName parameter is used by the integration tests;
    75  // specifically, each test case invocation generates a UUID which is used as the
    76  // name of the ProwJob CR. Then when the test pod is created, it is also named
    77  // with the same UUID. This makes checking for the creation of jobs and pods
    78  // very easy in the tests.
    79  func (c *Client) PublishMessage(ctx context.Context, msg PubSubMessageForSub, topicID string) error {
    80  	bytes, err := json.Marshal(msg.Data)
    81  	if err != nil {
    82  		return fmt.Errorf("failed to marshal: %v", err)
    83  	}
    84  
    85  	t := c.pubsubClient.Topic(topicID)
    86  	result := t.Publish(ctx, &pubsub.Message{Data: bytes, Attributes: msg.Attributes})
    87  
    88  	id, err := result.Get(ctx)
    89  	if err != nil {
    90  		return fmt.Errorf("failed to publish: %v", err)
    91  	}
    92  
    93  	logrus.Infof("successfully published message %v; msg ID: %v", string(bytes), id)
    94  
    95  	return nil
    96  }
    97  
    98  // CreateSubscription creates a Pub/Sub topic and a corresponding subscription.
    99  func (c *Client) CreateSubscription(ctx context.Context, projectID, topicID, subscriptionID string) error {
   100  	topic, err := c.pubsubClient.CreateTopic(ctx, topicID)
   101  	if err != nil {
   102  		return err
   103  	}
   104  
   105  	if _, err := c.pubsubClient.CreateSubscription(ctx, subscriptionID, pubsub.SubscriptionConfig{
   106  		Topic:            topic,
   107  		AckDeadline:      10 * time.Second,
   108  		ExpirationPolicy: 25 * time.Hour,
   109  	}); err != nil {
   110  		return err
   111  	}
   112  
   113  	return nil
   114  }