go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/cv/internal/common/pubsub/test_util.go (about)

     1  // Copyright 2021 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package pubsub
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"time"
    21  
    22  	"cloud.google.com/go/pubsub"
    23  	"cloud.google.com/go/pubsub/pstest"
    24  	"google.golang.org/api/option"
    25  	"google.golang.org/grpc"
    26  	"google.golang.org/grpc/credentials/insecure"
    27  
    28  	"go.chromium.org/luci/common/data/stringset"
    29  	"go.chromium.org/luci/common/errors"
    30  	"go.chromium.org/luci/common/logging"
    31  )
    32  
    33  // TestPSServer embeds a pubsub test server for testing PullingBatchProcessor.
    34  //
    35  // It also exposes a client that can be used to talk to the test server and the
    36  // fields needed to retrieve messages from a subscription.
    37  type TestPSServer struct {
    38  	*pstest.Server
    39  
    40  	Client *pubsub.Client
    41  	ProjID string
    42  	SubID  string
    43  
    44  	topicID, subName, topicName string
    45  	closers                     []func() error
    46  }
    47  
    48  // NewTestPSServer starts a test pubsub server, binds a client to it,
    49  // creates a default topic and subscription.
    50  func NewTestPSServer(ctx context.Context) (*TestPSServer, error) {
    51  	ret := &TestPSServer{}
    52  	var err error
    53  	defer func() {
    54  		if err != nil {
    55  			if closeErr := ret.Close(); closeErr != nil {
    56  				logging.Errorf(ctx, "Error closing mock pubsub server: %s", closeErr)
    57  			}
    58  		}
    59  	}()
    60  
    61  	ret.topicID, ret.ProjID, ret.SubID = "all-builds", "some-project", "some-sub-id"
    62  	ret.subName = fmt.Sprintf("projects/%s/subscriptions/%s", ret.ProjID, ret.SubID)
    63  	ret.topicName = fmt.Sprintf("projects/%s/topics/%s", ret.ProjID, ret.topicID)
    64  	ret.Server = pstest.NewServer()
    65  	conn, err := grpc.Dial(ret.Server.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  	ret.closers = append(ret.closers, conn.Close)
    70  	ret.Client, err = pubsub.NewClient(ctx, ret.ProjID, option.WithGRPCConn(conn))
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  	ret.closers = append(ret.closers, ret.Client.Close)
    75  	topic, err := ret.Client.CreateTopic(ctx, ret.topicID)
    76  	if err != nil {
    77  		return nil, err
    78  	}
    79  	_, err = ret.Client.CreateSubscription(ctx, ret.SubID, pubsub.SubscriptionConfig{
    80  		Topic:       topic,
    81  		AckDeadline: 600 * time.Second,
    82  	})
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	return ret, nil
    87  }
    88  
    89  // Close closes the embedded server and also the pubsub client and underlying
    90  // grpc connection.
    91  func (mps *TestPSServer) Close() error {
    92  	errs := &errors.MultiError{}
    93  	for i := len(mps.closers) - 1; i >= 0; i-- {
    94  		errs.MaybeAdd(mps.closers[i]())
    95  	}
    96  	errs.MaybeAdd(mps.Server.Close())
    97  	return errs.AsError()
    98  }
    99  
   100  // PublishTestMessages puts `n` distinct messages into the test servers default
   101  // topic.
   102  func (mps *TestPSServer) PublishTestMessages(n int) stringset.Set {
   103  	out := stringset.New(n)
   104  	for i := 0; i < n; i++ {
   105  		msgStr := fmt.Sprintf("msg%03d", i)
   106  		out.Add(msgStr)
   107  		// The call below returns the ID of the mock message, ignore it.
   108  		_ = mps.Publish(mps.topicName, []byte(msgStr), nil)
   109  	}
   110  	return out
   111  }