go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/quotabeta/admin_test.go (about)

     1  // Copyright 2022 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 quota
    16  
    17  import (
    18  	"context"
    19  	"strconv"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/alicebob/miniredis/v2"
    24  	"github.com/gomodule/redigo/redis"
    25  
    26  	"go.chromium.org/luci/common/clock/testclock"
    27  
    28  	pb "go.chromium.org/luci/server/quotabeta/proto"
    29  	"go.chromium.org/luci/server/quotabeta/quotaconfig"
    30  	"go.chromium.org/luci/server/redisconn"
    31  
    32  	. "github.com/smartystreets/goconvey/convey"
    33  	. "go.chromium.org/luci/common/testing/assertions"
    34  )
    35  
    36  func TestQuotaAdmin(t *testing.T) {
    37  	t.Parallel()
    38  
    39  	Convey("quotaAdmin", t, func() {
    40  		ctx, tc := testclock.UseTime(context.Background(), testclock.TestRecentTimeLocal)
    41  		now := strconv.FormatInt(tc.Now().Unix(), 10)
    42  		s, err := miniredis.Run()
    43  		So(err, ShouldBeNil)
    44  		defer s.Close()
    45  		ctx = redisconn.UsePool(ctx, &redis.Pool{
    46  			Dial: func() (redis.Conn, error) {
    47  				return redis.Dial("tcp", s.Addr())
    48  			},
    49  		})
    50  		conn, err := redisconn.Get(ctx)
    51  		So(err, ShouldBeNil)
    52  		_, err = conn.Do("HINCRBY", "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources", 3)
    53  		So(err, ShouldBeNil)
    54  		_, err = conn.Do("HINCRBY", "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated", tc.Now().Unix())
    55  		So(err, ShouldBeNil)
    56  		m, err := quotaconfig.NewMemory(ctx, []*pb.Policy{
    57  			{
    58  				Name:          "quota",
    59  				Resources:     2,
    60  				Replenishment: 1,
    61  			},
    62  			{
    63  				Name:          "quota/${user}",
    64  				Resources:     5,
    65  				Replenishment: 1,
    66  			},
    67  		})
    68  		So(err, ShouldBeNil)
    69  		ctx = Use(ctx, m)
    70  		srv := &quotaAdmin{}
    71  
    72  		Convey("Get", func() {
    73  			Convey("nil", func() {
    74  				rsp, err := srv.Get(ctx, nil)
    75  				So(err, ShouldErrLike, "policy is required")
    76  				So(rsp, ShouldBeNil)
    77  			})
    78  
    79  			Convey("empty", func() {
    80  				req := &pb.GetRequest{}
    81  				rsp, err := srv.Get(ctx, req)
    82  				So(err, ShouldErrLike, "policy is required")
    83  				So(rsp, ShouldBeNil)
    84  			})
    85  
    86  			Convey("not found", func() {
    87  				req := &pb.GetRequest{
    88  					Policy: "quota/user",
    89  				}
    90  
    91  				rsp, err := srv.Get(ctx, req)
    92  				So(err, ShouldErrLike, "not found")
    93  				So(rsp, ShouldBeNil)
    94  			})
    95  
    96  			Convey("user unspecified", func() {
    97  				req := &pb.GetRequest{
    98  					Policy: "quota/${user}",
    99  				}
   100  
   101  				rsp, err := srv.Get(ctx, req)
   102  				So(err, ShouldErrLike, "user not specified")
   103  				So(rsp, ShouldBeNil)
   104  			})
   105  
   106  			Convey("new", func() {
   107  				req := &pb.GetRequest{
   108  					Policy: "quota",
   109  				}
   110  
   111  				rsp, err := srv.Get(ctx, req)
   112  				So(err, ShouldBeNil)
   113  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   114  					Name:      "quota",
   115  					DbName:    "entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7",
   116  					Resources: 2,
   117  				})
   118  
   119  				// Ensure an entry for "quota" was not written to the database.
   120  				So(s.Keys(), ShouldResemble, []string{
   121  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   122  				})
   123  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources"), ShouldEqual, "3")
   124  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated"), ShouldEqual, now)
   125  			})
   126  
   127  			Convey("existing", func() {
   128  				req := &pb.GetRequest{
   129  					Policy: "quota/${user}",
   130  					User:   "user@example.com",
   131  				}
   132  
   133  				rsp, err := srv.Get(ctx, req)
   134  				So(err, ShouldBeNil)
   135  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   136  					Name:      "quota/user@example.com",
   137  					DbName:    "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   138  					Resources: 3,
   139  				})
   140  
   141  				So(s.Keys(), ShouldResemble, []string{
   142  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   143  				})
   144  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources"), ShouldEqual, "3")
   145  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated"), ShouldEqual, now)
   146  			})
   147  
   148  			Convey("replenish", func() {
   149  				ctx, _ := testclock.UseTime(ctx, testclock.TestRecentTimeLocal.Add(time.Second))
   150  				req := &pb.GetRequest{
   151  					Policy: "quota/${user}",
   152  					User:   "user@example.com",
   153  				}
   154  
   155  				rsp, err := srv.Get(ctx, req)
   156  				So(err, ShouldBeNil)
   157  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   158  					Name:      "quota/user@example.com",
   159  					DbName:    "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   160  					Resources: 4,
   161  				})
   162  
   163  				// Ensure replenishment was not written to the database.
   164  				So(s.Keys(), ShouldResemble, []string{
   165  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   166  				})
   167  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources"), ShouldEqual, "3")
   168  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated"), ShouldEqual, now)
   169  			})
   170  		})
   171  
   172  		Convey("Set", func() {
   173  			Convey("nil", func() {
   174  				rsp, err := srv.Set(ctx, nil)
   175  				So(err, ShouldErrLike, "policy is required")
   176  				So(rsp, ShouldBeNil)
   177  			})
   178  
   179  			Convey("empty", func() {
   180  				req := &pb.SetRequest{}
   181  				rsp, err := srv.Set(ctx, req)
   182  				So(err, ShouldErrLike, "policy is required")
   183  				So(rsp, ShouldBeNil)
   184  			})
   185  
   186  			Convey("negative", func() {
   187  				req := &pb.SetRequest{
   188  					Policy:    "quota",
   189  					Resources: -1,
   190  				}
   191  				rsp, err := srv.Set(ctx, req)
   192  				So(err, ShouldErrLike, "resources must not be negative")
   193  				So(rsp, ShouldBeNil)
   194  			})
   195  
   196  			Convey("not found", func() {
   197  				req := &pb.SetRequest{
   198  					Policy: "quota/user",
   199  				}
   200  
   201  				rsp, err := srv.Set(ctx, req)
   202  				So(err, ShouldErrLike, "not found")
   203  				So(rsp, ShouldBeNil)
   204  			})
   205  
   206  			Convey("user unspecified", func() {
   207  				req := &pb.SetRequest{
   208  					Policy: "quota/${user}",
   209  				}
   210  
   211  				rsp, err := srv.Set(ctx, req)
   212  				So(err, ShouldErrLike, "user not specified")
   213  				So(rsp, ShouldBeNil)
   214  			})
   215  
   216  			Convey("new", func() {
   217  				req := &pb.SetRequest{
   218  					Policy:    "quota",
   219  					Resources: 2,
   220  				}
   221  
   222  				rsp, err := srv.Set(ctx, req)
   223  				So(err, ShouldBeNil)
   224  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   225  					Name:      "quota",
   226  					DbName:    "entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7",
   227  					Resources: 2,
   228  				})
   229  
   230  				So(s.Keys(), ShouldResemble, []string{
   231  					"entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7",
   232  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   233  				})
   234  				So(s.HGet("entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7", "resources"), ShouldEqual, "2")
   235  				So(s.HGet("entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7", "updated"), ShouldEqual, now)
   236  			})
   237  
   238  			Convey("existing", func() {
   239  				req := &pb.SetRequest{
   240  					Policy:    "quota/${user}",
   241  					User:      "user@example.com",
   242  					Resources: 2,
   243  				}
   244  
   245  				rsp, err := srv.Set(ctx, req)
   246  				So(err, ShouldBeNil)
   247  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   248  					Name:      "quota/user@example.com",
   249  					DbName:    "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   250  					Resources: 2,
   251  				})
   252  
   253  				So(s.Keys(), ShouldResemble, []string{
   254  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   255  				})
   256  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources"), ShouldEqual, "2")
   257  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated"), ShouldEqual, now)
   258  			})
   259  
   260  			Convey("zero", func() {
   261  				req := &pb.SetRequest{
   262  					Policy: "quota/${user}",
   263  					User:   "user@example.com",
   264  				}
   265  
   266  				rsp, err := srv.Set(ctx, req)
   267  				So(err, ShouldBeNil)
   268  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   269  					Name:   "quota/user@example.com",
   270  					DbName: "entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   271  				})
   272  
   273  				So(s.Keys(), ShouldResemble, []string{
   274  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   275  				})
   276  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "resources"), ShouldEqual, "0")
   277  				So(s.HGet("entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d", "updated"), ShouldEqual, now)
   278  			})
   279  
   280  			Convey("excessive", func() {
   281  				req := &pb.SetRequest{
   282  					Policy:    "quota",
   283  					Resources: 10,
   284  				}
   285  
   286  				rsp, err := srv.Set(ctx, req)
   287  				So(err, ShouldBeNil)
   288  				So(rsp, ShouldResemble, &pb.QuotaEntry{
   289  					Name:      "quota",
   290  					DbName:    "entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7",
   291  					Resources: 2,
   292  				})
   293  
   294  				// Ensure resources were capped in the database.
   295  				So(s.Keys(), ShouldResemble, []string{
   296  					"entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7",
   297  					"entry:f20c860d2ea007ea2360c6ebe2d943acc8a531412c18ff3bd47ab1449988aa6d",
   298  				})
   299  				So(s.HGet("entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7", "resources"), ShouldEqual, "2")
   300  				So(s.HGet("entry:b878a6801d9a9e68b30ed63430bb5e0bddcd984a37a3ee385abc27ff031c7fe7", "updated"), ShouldEqual, now)
   301  			})
   302  		})
   303  	})
   304  }