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

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