go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/swarming/server/pubsub/pubsub_test.go (about)

     1  // Copyright 2023 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  	"bytes"
    19  	"context"
    20  	"encoding/json"
    21  	"net/http"
    22  	"net/http/httptest"
    23  	"net/url"
    24  	"testing"
    25  	"time"
    26  
    27  	"google.golang.org/protobuf/proto"
    28  	"google.golang.org/protobuf/types/known/timestamppb"
    29  
    30  	"go.chromium.org/luci/common/errors"
    31  	"go.chromium.org/luci/common/retry/transient"
    32  	"go.chromium.org/luci/server/auth"
    33  	"go.chromium.org/luci/server/auth/authtest"
    34  	"go.chromium.org/luci/server/router"
    35  
    36  	. "github.com/smartystreets/goconvey/convey"
    37  	. "go.chromium.org/luci/common/testing/assertions"
    38  )
    39  
    40  func TestPubSubHandler(t *testing.T) {
    41  	t.Parallel()
    42  
    43  	Convey("With mocks", t, func() {
    44  		const expectedPusherID = "user:push@example.com"
    45  		var (
    46  			testTime = time.Unix(1689000000, 0)
    47  			testMsg  = &timestamppb.Timestamp{Seconds: 123456}
    48  		)
    49  
    50  		var body pushRequestBody
    51  		body.Subscription = "sub"
    52  		body.Message.Attributes = map[string]string{"a1": "v1", "a2": "v2"}
    53  		body.Message.MessageID = "msg"
    54  		body.Message.PublishTime = testTime
    55  		body.Message.Data, _ = proto.Marshal(testMsg)
    56  
    57  		goodBlob, _ := json.Marshal(&body)
    58  
    59  		ctx := auth.WithState(context.Background(),
    60  			&authtest.FakeState{
    61  				Identity: expectedPusherID,
    62  			},
    63  		)
    64  
    65  		call := func(
    66  			ctx context.Context,
    67  			path string,
    68  			body []byte,
    69  			cb func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error,
    70  		) (statusCode int) {
    71  			rr := httptest.NewRecorder()
    72  			rc := &router.Context{
    73  				Request: httptest.NewRequest("POST", path, bytes.NewReader(body)).WithContext(ctx),
    74  				Writer:  rr,
    75  			}
    76  			handler(rc, expectedPusherID, cb)
    77  			return rr.Code
    78  		}
    79  
    80  		Convey("Success", func() {
    81  			var gotMsg *timestamppb.Timestamp
    82  			var gotMD *Metadata
    83  			resp := call(ctx, "/p?a=1&b=2", goodBlob, func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
    84  				gotMsg = msg
    85  				gotMD = md
    86  				return nil
    87  			})
    88  
    89  			So(resp, ShouldEqual, http.StatusOK)
    90  			So(gotMsg, ShouldResembleProto, testMsg)
    91  			So(gotMD.Subscription, ShouldEqual, "sub")
    92  			So(gotMD.MessageID, ShouldEqual, "msg")
    93  			So(gotMD.PublishTime, ShouldEqual, testTime)
    94  			So(gotMD.Attributes, ShouldResemble, body.Message.Attributes)
    95  			So(gotMD.Query, ShouldResemble, url.Values{"a": {"1"}, "b": {"2"}})
    96  		})
    97  
    98  		Convey("Transient error", func() {
    99  			resp := call(ctx, "/p", goodBlob, func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
   100  				return errors.New("boo", transient.Tag)
   101  			})
   102  			So(resp, ShouldEqual, http.StatusInternalServerError)
   103  		})
   104  
   105  		Convey("Fatal error", func() {
   106  			resp := call(ctx, "/p", goodBlob, func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
   107  				return errors.New("boo")
   108  			})
   109  			So(resp, ShouldEqual, http.StatusAccepted)
   110  		})
   111  
   112  		Convey("Wrong caller ID", func() {
   113  			ctx := auth.WithState(context.Background(),
   114  				&authtest.FakeState{
   115  					Identity: "user:wrong@example.com",
   116  				},
   117  			)
   118  			resp := call(ctx, "/p", goodBlob, func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
   119  				return nil
   120  			})
   121  			So(resp, ShouldEqual, http.StatusForbidden)
   122  		})
   123  
   124  		Convey("Bad wrapper", func() {
   125  			resp := call(ctx, "/p", []byte("not json"), func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
   126  				return nil
   127  			})
   128  			So(resp, ShouldEqual, http.StatusBadRequest)
   129  		})
   130  
   131  		Convey("Bad payload", func() {
   132  			var body pushRequestBody
   133  			body.Message.Data = []byte("bad proto")
   134  			blob, _ := json.Marshal(&body)
   135  			resp := call(ctx, "/p", blob, func(ctx context.Context, msg *timestamppb.Timestamp, md *Metadata) error {
   136  				return nil
   137  			})
   138  			So(resp, ShouldEqual, http.StatusBadRequest)
   139  		})
   140  	})
   141  }