services/crunch-dispatch-local/crunch-dispatch-local_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "bytes"
   9         "context"
  10         "io"
  11         "net/http"
  12         "net/http/httptest"
  13         "os"
  14         "os/exec"
  15         "regexp"
  16         "testing"
  17         "time"
  18
  19         "git.arvados.org/arvados.git/sdk/go/arvados"
  20         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  21         "git.arvados.org/arvados.git/sdk/go/arvadostest"
  22         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  23         "git.arvados.org/arvados.git/sdk/go/dispatch"
  24         . "gopkg.in/check.v1"
  25 )
  26
  27 // Gocheck boilerplate
  28 func Test(t *testing.T) {
  29         TestingT(t)
  30 }
  31
  32 var _ = Suite(&TestSuite{})
  33 var _ = Suite(&MockArvadosServerSuite{})
  34
  35 type TestSuite struct{}
  36 type MockArvadosServerSuite struct{}
  37
  38 var initialArgs []string
  39
  40 func (s *TestSuite) SetUpSuite(c *C) {
  41         initialArgs = os.Args
  42         runningCmds = make(map[string]*exec.Cmd)
  43 }
  44
  45 func (s *TestSuite) SetUpTest(c *C) {
  46         arvadostest.ResetDB(c)
  47         arvadostest.ResetEnv()
  48         args := []string{"crunch-dispatch-local"}
  49         os.Args = args
  50 }
  51
  52 func (s *TestSuite) TearDownTest(c *C) {
  53         arvadostest.ResetEnv()
  54         arvadostest.ResetDB(c)
  55         os.Args = initialArgs
  56 }
  57
  58 func (s *MockArvadosServerSuite) TearDownTest(c *C) {
  59         arvadostest.ResetEnv()
  60 }
  61
  62 func (s *TestSuite) TestIntegration(c *C) {
  63         arv, err := arvadosclient.MakeArvadosClient()
  64         c.Assert(err, IsNil)
  65
  66         echo := "echo"
  67         crunchRunCommand = echo
  68
  69         ctx, cancel := context.WithCancel(ctxlog.Context(context.Background(), ctxlog.TestLogger(c)))
  70         dispatcher := dispatch.Dispatcher{
  71                 Arv:        arv,
  72                 PollPeriod: time.Second,
  73         }
  74
  75         startCmd := func(container arvados.Container, cmd *exec.Cmd) error {
  76                 dispatcher.UpdateState(container.UUID, "Running")
  77                 dispatcher.UpdateState(container.UUID, "Complete")
  78                 return cmd.Start()
  79         }
  80
  81         cl := arvados.Cluster{Containers: arvados.ContainersConfig{RuntimeEngine: "docker"}}
  82
  83         dispatcher.RunContainer = func(d *dispatch.Dispatcher, c arvados.Container, s <-chan arvados.Container) error {
  84                 defer cancel()
  85                 return (&LocalRun{startCmd, make(chan bool, 8), ctx, &cl}).run(d, c, s)
  86         }
  87
  88         err = dispatcher.Run(ctx)
  89         c.Assert(err, Equals, context.Canceled)
  90
  91         // Wait for all running crunch jobs to complete / terminate
  92         waitGroup.Wait()
  93
  94         // There should be no queued containers now
  95         params := arvadosclient.Dict{
  96                 "filters": [][]string{{"state", "=", "Queued"}},
  97         }
  98         var containers arvados.ContainerList
  99         err = arv.List("containers", params, &containers)
 100         c.Check(err, IsNil)
 101         c.Assert(len(containers.Items), Equals, 0)
 102
 103         // Previously "Queued" container should now be in "Complete" state
 104         var container arvados.Container
 105         err = arv.Get("containers", "zzzzz-dz642-queuedcontainer", nil, &container)
 106         c.Check(err, IsNil)
 107         c.Check(string(container.State), Equals, "Complete")
 108 }
 109
 110 func (s *MockArvadosServerSuite) Test_APIErrorGettingContainers(c *C) {
 111         apiStubResponses := make(map[string]arvadostest.StubResponse)
 112         apiStubResponses["/arvados/v1/containers"] = arvadostest.StubResponse{500, string(`{}`)}
 113
 114         testWithServerStub(c, apiStubResponses, "echo", "error getting count of containers")
 115 }
 116
 117 func (s *MockArvadosServerSuite) Test_APIErrorUpdatingContainerState(c *C) {
 118         apiStubResponses := make(map[string]arvadostest.StubResponse)
 119         apiStubResponses["/arvados/v1/containers"] =
 120                 arvadostest.StubResponse{200, string(`{"items_available":1, "items":[{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx1","State":"Queued","Priority":1}]}`)}
 121         apiStubResponses["/arvados/v1/containers/zzzzz-dz642-xxxxxxxxxxxxxx1"] =
 122                 arvadostest.StubResponse{500, string(`{}`)}
 123
 124         testWithServerStub(c, apiStubResponses, "echo", "error locking container zzzzz-dz642-xxxxxxxxxxxxxx1")
 125 }
 126
 127 func (s *MockArvadosServerSuite) Test_ContainerStillInRunningAfterRun(c *C) {
 128         apiStubResponses := make(map[string]arvadostest.StubResponse)
 129         apiStubResponses["/arvados/v1/containers"] =
 130                 arvadostest.StubResponse{200, string(`{"items_available":1, "items":[{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx2","State":"Queued","Priority":1}]}`)}
 131         apiStubResponses["/arvados/v1/containers/zzzzz-dz642-xxxxxxxxxxxxxx2/lock"] =
 132                 arvadostest.StubResponse{200, string(`{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx2", "state":"Locked", "priority":1, "locked_by_uuid": "` + arvadostest.Dispatch1AuthUUID + `"}`)}
 133         apiStubResponses["/arvados/v1/containers/zzzzz-dz642-xxxxxxxxxxxxxx2"] =
 134                 arvadostest.StubResponse{200, string(`{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx2", "state":"Running", "priority":1, "locked_by_uuid": "` + arvadostest.Dispatch1AuthUUID + `"}`)}
 135
 136         testWithServerStub(c, apiStubResponses, "echo",
 137                 `after \\"echo\\" process termination, container state for zzzzz-dz642-xxxxxxxxxxxxxx2 is \\"Running\\"; updating it to \\"Cancelled\\"`)
 138 }
 139
 140 func (s *MockArvadosServerSuite) Test_ErrorRunningContainer(c *C) {
 141         apiStubResponses := make(map[string]arvadostest.StubResponse)
 142         apiStubResponses["/arvados/v1/containers"] =
 143                 arvadostest.StubResponse{200, string(`{"items_available":1, "items":[{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx3","State":"Queued","Priority":1}]}`)}
 144
 145         apiStubResponses["/arvados/v1/containers/zzzzz-dz642-xxxxxxxxxxxxxx3/lock"] =
 146                 arvadostest.StubResponse{200, string(`{"uuid":"zzzzz-dz642-xxxxxxxxxxxxxx3", "state":"Locked", "priority":1}`)}
 147
 148         testWithServerStub(c, apiStubResponses, "nosuchcommand", `error starting \\"nosuchcommand\\" for zzzzz-dz642-xxxxxxxxxxxxxx3`)
 149 }
 150
 151 func testWithServerStub(c *C, apiStubResponses map[string]arvadostest.StubResponse, crunchCmd string, expected string) {
 152         apiStubResponses["/arvados/v1/api_client_authorizations/current"] =
 153                 arvadostest.StubResponse{200, string(`{"uuid": "` + arvadostest.Dispatch1AuthUUID + `", "api_token": "xyz"}`)}
 154
 155         apiStub := arvadostest.ServerStub{apiStubResponses}
 156
 157         api := httptest.NewServer(&apiStub)
 158         defer api.Close()
 159
 160         arv := &arvadosclient.ArvadosClient{
 161                 Scheme:    "http",
 162                 ApiServer: api.URL[7:],
 163                 ApiToken:  "abc123",
 164                 Client:    &http.Client{Transport: &http.Transport{}},
 165                 Retries:   0,
 166         }
 167
 168         buf := bytes.NewBuffer(nil)
 169         logger := ctxlog.TestLogger(c)
 170         logger.SetOutput(io.MultiWriter(buf, logger.Out))
 171
 172         crunchRunCommand = crunchCmd
 173
 174         ctx, cancel := context.WithCancel(ctxlog.Context(context.Background(), logger))
 175         defer cancel()
 176         dispatcher := dispatch.Dispatcher{
 177                 Logger:     logger,
 178                 Arv:        arv,
 179                 PollPeriod: time.Second,
 180         }
 181
 182         startCmd := func(container arvados.Container, cmd *exec.Cmd) error {
 183                 dispatcher.UpdateState(container.UUID, "Running")
 184                 dispatcher.UpdateState(container.UUID, "Complete")
 185                 return cmd.Start()
 186         }
 187
 188         cl := arvados.Cluster{Containers: arvados.ContainersConfig{RuntimeEngine: "docker"}}
 189
 190         dispatcher.RunContainer = func(d *dispatch.Dispatcher, c arvados.Container, s <-chan arvados.Container) error {
 191                 defer cancel()
 192                 return (&LocalRun{startCmd, make(chan bool, 8), ctx, &cl}).run(d, c, s)
 193         }
 194
 195         re := regexp.MustCompile(`(?ms).*` + expected + `.*`)
 196         go func() {
 197                 for i := 0; i < 80 && !re.MatchString(buf.String()); i++ {
 198                         time.Sleep(100 * time.Millisecond)
 199                 }
 200                 cancel()
 201         }()
 202
 203         err := dispatcher.Run(ctx)
 204         c.Assert(err, Equals, context.Canceled)
 205         c.Check(buf.String(), Matches, `(?ms).*`+expected+`.*`)
 206
 207         c.Logf("test finished, waiting for running crunch jobs to complete / terminate")
 208         waitGroup.Wait()
 209 }