go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/appengine/gaemiddleware/appengine_test.go (about)

     1  // Copyright 2015 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 gaemiddleware
    16  
    17  import (
    18  	"context"
    19  	"net/http"
    20  	"net/http/httptest"
    21  	"testing"
    22  
    23  	. "github.com/smartystreets/goconvey/convey"
    24  	"go.chromium.org/luci/appengine/gaetesting"
    25  	"go.chromium.org/luci/server/router"
    26  )
    27  
    28  func init() {
    29  	// disable this so that we can actually check the logic in these middlewares
    30  	devAppserverBypassFn = func(context.Context) bool { return false }
    31  }
    32  
    33  func TestRequireCron(t *testing.T) {
    34  	t.Parallel()
    35  
    36  	Convey("Test RequireCron", t, func() {
    37  		hit := false
    38  		f := func(c *router.Context) {
    39  			hit = true
    40  			c.Writer.Write([]byte("ok"))
    41  		}
    42  
    43  		Convey("from non-cron fails", func() {
    44  			rec := httptest.NewRecorder()
    45  			c := &router.Context{
    46  				Writer:  rec,
    47  				Request: (&http.Request{}).WithContext(gaetesting.TestingContext()),
    48  			}
    49  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireCron), f)
    50  			So(hit, ShouldBeFalse)
    51  			So(rec.Body.String(), ShouldEqual, "error: must be run from cron\n")
    52  			So(rec.Code, ShouldEqual, http.StatusForbidden)
    53  		})
    54  
    55  		Convey("from cron succeeds", func() {
    56  			rec := httptest.NewRecorder()
    57  			c := &router.Context{
    58  				Writer: rec,
    59  				Request: (&http.Request{
    60  					Header: http.Header{
    61  						http.CanonicalHeaderKey("x-appengine-cron"): []string{"true"},
    62  					},
    63  				}).WithContext(gaetesting.TestingContext()),
    64  			}
    65  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireCron), f)
    66  			So(hit, ShouldBeTrue)
    67  			So(rec.Body.String(), ShouldEqual, "ok")
    68  			So(rec.Code, ShouldEqual, http.StatusOK)
    69  		})
    70  	})
    71  }
    72  
    73  func TestRequireTQ(t *testing.T) {
    74  	t.Parallel()
    75  
    76  	Convey("Test RequireTQ", t, func() {
    77  		hit := false
    78  		f := func(c *router.Context) {
    79  			hit = true
    80  			c.Writer.Write([]byte("ok"))
    81  		}
    82  
    83  		Convey("from non-tq fails (wat)", func() {
    84  			rec := httptest.NewRecorder()
    85  			c := &router.Context{
    86  				Writer:  rec,
    87  				Request: (&http.Request{}).WithContext(gaetesting.TestingContext()),
    88  			}
    89  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireTaskQueue("wat")), f)
    90  			So(hit, ShouldBeFalse)
    91  			So(rec.Body.String(), ShouldEqual, "error: must be run from the correct taskqueue\n")
    92  			So(rec.Code, ShouldEqual, http.StatusForbidden)
    93  		})
    94  
    95  		Convey("from non-tq fails", func() {
    96  			rec := httptest.NewRecorder()
    97  			c := &router.Context{
    98  				Writer:  rec,
    99  				Request: (&http.Request{}).WithContext(gaetesting.TestingContext()),
   100  			}
   101  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireTaskQueue("")), f)
   102  			So(hit, ShouldBeFalse)
   103  			So(rec.Body.String(), ShouldEqual, "error: must be run from the correct taskqueue\n")
   104  			So(rec.Code, ShouldEqual, http.StatusForbidden)
   105  		})
   106  
   107  		Convey("from wrong tq fails (wat)", func() {
   108  			rec := httptest.NewRecorder()
   109  			c := &router.Context{
   110  				Writer: rec,
   111  				Request: (&http.Request{
   112  					Header: http.Header{
   113  						http.CanonicalHeaderKey("x-appengine-queuename"): []string{"else"},
   114  					},
   115  				}).WithContext(gaetesting.TestingContext()),
   116  			}
   117  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireTaskQueue("wat")), f)
   118  			So(hit, ShouldBeFalse)
   119  			So(rec.Body.String(), ShouldEqual, "error: must be run from the correct taskqueue\n")
   120  			So(rec.Code, ShouldEqual, http.StatusForbidden)
   121  		})
   122  
   123  		Convey("from right tq succeeds (wat)", func() {
   124  			rec := httptest.NewRecorder()
   125  			c := &router.Context{
   126  				Writer: rec,
   127  				Request: (&http.Request{
   128  					Header: http.Header{
   129  						http.CanonicalHeaderKey("x-appengine-queuename"): []string{"wat"},
   130  					},
   131  				}).WithContext(gaetesting.TestingContext()),
   132  			}
   133  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireTaskQueue("wat")), f)
   134  			So(hit, ShouldBeTrue)
   135  			So(rec.Body.String(), ShouldEqual, "ok")
   136  			So(rec.Code, ShouldEqual, http.StatusOK)
   137  		})
   138  
   139  		Convey("from any tq succeeds", func() {
   140  			rec := httptest.NewRecorder()
   141  			c := &router.Context{
   142  				Writer: rec,
   143  				Request: (&http.Request{
   144  					Header: http.Header{
   145  						http.CanonicalHeaderKey("x-appengine-queuename"): []string{"wat"},
   146  					},
   147  				}).WithContext(gaetesting.TestingContext()),
   148  			}
   149  			router.RunMiddleware(c, router.NewMiddlewareChain(RequireTaskQueue("")), f)
   150  			So(hit, ShouldBeTrue)
   151  			So(rec.Body.String(), ShouldEqual, "ok")
   152  			So(rec.Code, ShouldEqual, http.StatusOK)
   153  		})
   154  	})
   155  }