go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gce/appengine/backend/cron_test.go (about)

     1  // Copyright 2018 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 backend
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"net/http"
    21  	"testing"
    22  
    23  	"google.golang.org/api/compute/v1"
    24  	"google.golang.org/api/option"
    25  
    26  	"go.chromium.org/luci/appengine/tq"
    27  	"go.chromium.org/luci/appengine/tq/tqtesting"
    28  	"go.chromium.org/luci/gae/impl/memory"
    29  	"go.chromium.org/luci/gae/service/datastore"
    30  	"go.chromium.org/luci/gce/api/config/v1"
    31  	"go.chromium.org/luci/gce/api/projects/v1"
    32  	"go.chromium.org/luci/gce/api/tasks/v1"
    33  	"go.chromium.org/luci/gce/appengine/model"
    34  	"go.chromium.org/luci/gce/appengine/testing/roundtripper"
    35  
    36  	. "github.com/smartystreets/goconvey/convey"
    37  	. "go.chromium.org/luci/common/testing/assertions"
    38  )
    39  
    40  func TestCron(t *testing.T) {
    41  	t.Parallel()
    42  
    43  	Convey("cron", t, func() {
    44  		dsp := &tq.Dispatcher{}
    45  		registerTasks(dsp)
    46  		rt := &roundtripper.JSONRoundTripper{}
    47  		c := context.Background()
    48  		gce, err := compute.NewService(c, option.WithHTTPClient(&http.Client{Transport: rt}))
    49  		So(err, ShouldBeNil)
    50  		c = withCompute(withDispatcher(memory.Use(c), dsp), ComputeService{Stable: gce})
    51  		datastore.GetTestable(c).Consistent(true)
    52  		tqt := tqtesting.GetTestable(c, dsp)
    53  		tqt.CreateQueues()
    54  
    55  		Convey("countTasks", func() {
    56  			dsp := &tq.Dispatcher{}
    57  			c = withDispatcher(c, dsp)
    58  			q := datastore.NewQuery("TaskCount")
    59  
    60  			Convey("none", func() {
    61  				So(countTasks(c), ShouldBeNil)
    62  				var k []*datastore.Key
    63  				So(datastore.GetAll(c, q, &k), ShouldBeNil)
    64  				So(k, ShouldBeEmpty)
    65  			})
    66  
    67  			Convey("many", func() {
    68  				dsp.RegisterTask(&tasks.CountVMs{}, countVMs, countVMsQueue, nil)
    69  				dsp.RegisterTask(&tasks.ManageBot{}, manageBot, manageBotQueue, nil)
    70  				So(countTasks(c), ShouldBeNil)
    71  				var k []*datastore.Key
    72  				So(datastore.GetAll(c, q, &k), ShouldBeNil)
    73  				So(k, ShouldHaveLength, 2)
    74  			})
    75  		})
    76  
    77  		Convey("countVMsAsync", func() {
    78  			Convey("none", func() {
    79  				So(countVMsAsync(c), ShouldBeNil)
    80  				So(tqt.GetScheduledTasks(), ShouldBeEmpty)
    81  			})
    82  
    83  			Convey("one", func() {
    84  				datastore.Put(c, &model.Config{
    85  					ID: "id",
    86  				})
    87  				So(countVMsAsync(c), ShouldBeNil)
    88  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
    89  			})
    90  		})
    91  
    92  		Convey("createInstancesAsync", func() {
    93  			Convey("none", func() {
    94  				Convey("zero", func() {
    95  					So(createInstancesAsync(c), ShouldBeNil)
    96  					So(tqt.GetScheduledTasks(), ShouldBeEmpty)
    97  				})
    98  
    99  				Convey("exists", func() {
   100  					datastore.Put(c, &model.VM{
   101  						ID:  "id",
   102  						URL: "url",
   103  					})
   104  					So(createInstancesAsync(c), ShouldBeNil)
   105  					So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   106  				})
   107  			})
   108  
   109  			Convey("one", func() {
   110  				datastore.Put(c, &model.VM{
   111  					ID: "id",
   112  					Attributes: config.VM{
   113  						Zone: "zone",
   114  					},
   115  				})
   116  				So(createInstancesAsync(c), ShouldBeNil)
   117  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   118  			})
   119  		})
   120  
   121  		Convey("expandConfigsAsync", func() {
   122  			Convey("none", func() {
   123  				So(expandConfigsAsync(c), ShouldBeNil)
   124  				So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   125  			})
   126  
   127  			Convey("one", func() {
   128  				datastore.Put(c, &model.Config{
   129  					ID: "id",
   130  				})
   131  				So(expandConfigsAsync(c), ShouldBeNil)
   132  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   133  			})
   134  
   135  			Convey("many", func() {
   136  				for i := 0; i < 100; i++ {
   137  					datastore.Put(c, &model.Config{
   138  						ID: fmt.Sprintf("id-%d", i),
   139  					})
   140  				}
   141  				So(expandConfigsAsync(c), ShouldBeNil)
   142  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 100)
   143  			})
   144  		})
   145  
   146  		Convey("manageBotsAsync", func() {
   147  			Convey("none", func() {
   148  				Convey("missing", func() {
   149  					So(manageBotsAsync(c), ShouldBeNil)
   150  					So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   151  				})
   152  
   153  				Convey("url", func() {
   154  					datastore.Put(c, &model.VM{
   155  						ID: "id",
   156  					})
   157  					So(manageBotsAsync(c), ShouldBeNil)
   158  					So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   159  				})
   160  			})
   161  
   162  			Convey("one", func() {
   163  				datastore.Put(c, &model.VM{
   164  					ID:  "id",
   165  					URL: "url",
   166  				})
   167  				So(manageBotsAsync(c), ShouldBeNil)
   168  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   169  			})
   170  		})
   171  
   172  		Convey("drainVMsAsync", func() {
   173  			Convey("none", func() {
   174  				// No VMs will be drained as config is set for 2 VMs
   175  				So(datastore.Put(c, &model.VM{
   176  					ID:     "id-0",
   177  					Index:  0,
   178  					Config: "id",
   179  				}), ShouldBeNil)
   180  				So(datastore.Put(c, &model.VM{
   181  					ID:     "id-1",
   182  					Index:  1,
   183  					Config: "id",
   184  				}), ShouldBeNil)
   185  				So(datastore.Put(c, &model.Config{
   186  					ID: "id",
   187  					Config: &config.Config{
   188  						Prefix:        "id",
   189  						CurrentAmount: 2,
   190  					},
   191  				}), ShouldBeNil)
   192  				Convey("missing", func() {
   193  					So(drainVMsAsync(c), ShouldBeNil)
   194  					So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   195  				})
   196  			})
   197  
   198  			Convey("one", func() {
   199  				// 1 VM will be drained as config is set for 1 VM
   200  				So(datastore.Put(c, &model.VM{
   201  					ID:     "id-0",
   202  					Index:  0,
   203  					Config: "id",
   204  				}), ShouldBeNil)
   205  				So(datastore.Put(c, &model.VM{
   206  					ID:     "id-1",
   207  					Index:  1,
   208  					Config: "id",
   209  				}), ShouldBeNil)
   210  				So(datastore.Put(c, &model.Config{
   211  					ID: "id",
   212  					Config: &config.Config{
   213  						Prefix:        "id",
   214  						CurrentAmount: 1,
   215  					},
   216  				}), ShouldBeNil)
   217  				Convey("missing", func() {
   218  					So(drainVMsAsync(c), ShouldBeNil)
   219  					So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   220  				})
   221  			})
   222  
   223  			Convey("all", func() {
   224  				// 2 VMs will be drained as config is set for 0 VM
   225  				So(datastore.Put(c, &model.VM{
   226  					ID:     "id-0",
   227  					Index:  0,
   228  					Config: "id",
   229  				}), ShouldBeNil)
   230  				So(datastore.Put(c, &model.VM{
   231  					ID:     "id-1",
   232  					Index:  1,
   233  					Config: "id",
   234  				}), ShouldBeNil)
   235  				So(datastore.Put(c, &model.Config{
   236  					ID: "id",
   237  					Config: &config.Config{
   238  						Prefix:        "id",
   239  						CurrentAmount: 0,
   240  					},
   241  				}), ShouldBeNil)
   242  				Convey("missing", func() {
   243  					So(drainVMsAsync(c), ShouldBeNil)
   244  					So(tqt.GetScheduledTasks(), ShouldHaveLength, 2)
   245  				})
   246  			})
   247  		})
   248  
   249  		Convey("payloadFactory", func() {
   250  			f := payloadFactory(&tasks.CountVMs{})
   251  			p := f("id")
   252  			So(p, ShouldResemble, &tasks.CountVMs{
   253  				Id: "id",
   254  			})
   255  		})
   256  
   257  		Convey("reportQuotasAsync", func() {
   258  			Convey("none", func() {
   259  				So(reportQuotasAsync(c), ShouldBeNil)
   260  				So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   261  			})
   262  
   263  			Convey("one", func() {
   264  				datastore.Put(c, &model.Project{
   265  					ID: "id",
   266  				})
   267  				So(reportQuotasAsync(c), ShouldBeNil)
   268  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   269  			})
   270  		})
   271  
   272  		Convey("auditInstances", func() {
   273  			Convey("none", func() {
   274  				So(auditInstances(c), ShouldBeNil)
   275  				So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   276  			})
   277  
   278  			Convey("zero", func() {
   279  				count := 0
   280  				rt.Handler = func(req any) (int, any) {
   281  					switch count {
   282  					case 0:
   283  						count += 1
   284  						return http.StatusOK, &compute.ZoneList{
   285  							Items: []*compute.Zone{},
   286  						}
   287  					default:
   288  						count += 1
   289  						return http.StatusInternalServerError, nil
   290  					}
   291  				}
   292  				err := datastore.Put(c, &model.Project{
   293  					ID: "id",
   294  					Config: &projects.Config{
   295  						Project: "gnu-hurd",
   296  					},
   297  				})
   298  				So(err, ShouldBeNil)
   299  				So(auditInstances(c), ShouldBeNil)
   300  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 0)
   301  			})
   302  
   303  			Convey("one", func() {
   304  				count := 0
   305  				rt.Handler = func(req any) (int, any) {
   306  					switch count {
   307  					case 0:
   308  						count += 1
   309  						return http.StatusOK, &compute.ZoneList{
   310  							Items: []*compute.Zone{{
   311  								Name: "us-mex-1",
   312  							}, {
   313  								Name: "us-num-1",
   314  							}},
   315  						}
   316  					default:
   317  						count += 1
   318  						return http.StatusInternalServerError, nil
   319  					}
   320  				}
   321  				err := datastore.Put(c, &model.Project{
   322  					ID: "id",
   323  					Config: &projects.Config{
   324  						Project: "gnu-hurd",
   325  					},
   326  				})
   327  				So(err, ShouldBeNil)
   328  				So(auditInstances(c), ShouldBeNil)
   329  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 2)
   330  			})
   331  
   332  			Convey("two", func() {
   333  				count := 0
   334  				rt.Handler = func(req any) (int, any) {
   335  					switch count {
   336  					case 0:
   337  						count += 1
   338  						return http.StatusOK, &compute.ZoneList{
   339  							Items: []*compute.Zone{{
   340  								Name: "us-mex-1",
   341  							}, {
   342  								Name: "us-num-1",
   343  							}},
   344  						}
   345  					case 1:
   346  						count += 1
   347  						return http.StatusOK, &compute.ZoneList{
   348  							Items: []*compute.Zone{{
   349  								Name: "us-can-2",
   350  							}},
   351  						}
   352  					default:
   353  						count += 1
   354  						return http.StatusInternalServerError, nil
   355  					}
   356  				}
   357  				err := datastore.Put(c, &model.Project{
   358  					ID: "id",
   359  					Config: &projects.Config{
   360  						Project: "gnu-hurd",
   361  					},
   362  				})
   363  				So(err, ShouldBeNil)
   364  				err = datastore.Put(c, &model.Project{
   365  					ID: "id2",
   366  					Config: &projects.Config{
   367  						Project: "libreboot",
   368  					},
   369  				})
   370  				So(err, ShouldBeNil)
   371  				So(auditInstances(c), ShouldBeNil)
   372  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 3)
   373  			})
   374  		})
   375  
   376  		Convey("trigger", func() {
   377  			Convey("none", func() {
   378  				So(trigger(c, &tasks.ManageBot{}, datastore.NewQuery(model.VMKind)), ShouldBeNil)
   379  				So(tqt.GetScheduledTasks(), ShouldBeEmpty)
   380  			})
   381  
   382  			Convey("one", func() {
   383  				datastore.Put(c, &model.VM{
   384  					ID: "id",
   385  				})
   386  				So(trigger(c, &tasks.ManageBot{}, datastore.NewQuery(model.VMKind)), ShouldBeNil)
   387  				So(tqt.GetScheduledTasks(), ShouldHaveLength, 1)
   388  				So(tqt.GetScheduledTasks()[0].Payload, ShouldResembleProto, &tasks.ManageBot{
   389  					Id: "id",
   390  				})
   391  			})
   392  		})
   393  	})
   394  }