services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go

   1 package main
   2
   3 import (
   4         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
   5         "git.curoverse.com/arvados.git/sdk/go/arvadostest"
   6
   7         "io/ioutil"
   8         "log"
   9         "net/http"
  10         "net/http/httptest"
  11         "os"
  12         "os/exec"
  13         "strings"
  14         "syscall"
  15         "testing"
  16         "time"
  17
  18         . "gopkg.in/check.v1"
  19 )
  20
  21 // Gocheck boilerplate
  22 func Test(t *testing.T) {
  23         TestingT(t)
  24 }
  25
  26 var _ = Suite(&TestSuite{})
  27 var _ = Suite(&MockArvadosServerSuite{})
  28
  29 type TestSuite struct{}
  30 type MockArvadosServerSuite struct{}
  31
  32 var initialArgs []string
  33
  34 func (s *TestSuite) SetUpSuite(c *C) {
  35         initialArgs = os.Args
  36         arvadostest.StartAPI()
  37 }
  38
  39 func (s *TestSuite) TearDownSuite(c *C) {
  40         arvadostest.StopAPI()
  41 }
  42
  43 func (s *TestSuite) SetUpTest(c *C) {
  44         args := []string{"crunch-dispatch-slurm"}
  45         os.Args = args
  46
  47         var err error
  48         arv, err = arvadosclient.MakeArvadosClient()
  49         if err != nil {
  50                 c.Fatalf("Error making arvados client: %s", err)
  51         }
  52 }
  53
  54 func (s *TestSuite) TearDownTest(c *C) {
  55         arvadostest.ResetEnv()
  56         os.Args = initialArgs
  57 }
  58
  59 func (s *MockArvadosServerSuite) TearDownTest(c *C) {
  60         arvadostest.ResetEnv()
  61 }
  62
  63 func (s *TestSuite) Test_doMain(c *C) {
  64         args := []string{"-poll-interval", "2", "-container-priority-poll-interval", "1", "-crunch-run-command", "echo"}
  65         os.Args = append(os.Args, args...)
  66
  67         var sbatchCmdLine []string
  68         var striggerCmdLine []string
  69
  70         // Override sbatchCmd
  71         defer func(orig func(string) *exec.Cmd) {
  72                 sbatchCmd = orig
  73         }(sbatchCmd)
  74         sbatchCmd = func(uuid string) *exec.Cmd {
  75                 sbatchCmdLine = sbatchFunc(uuid).Args
  76                 return exec.Command("echo", uuid)
  77         }
  78
  79         // Override striggerCmd
  80         defer func(orig func(jobid, containerUUID, finishCommand,
  81                 apiHost, apiToken, apiInsecure string) *exec.Cmd) {
  82                 striggerCmd = orig
  83         }(striggerCmd)
  84         striggerCmd = func(jobid, containerUUID, finishCommand, apiHost, apiToken, apiInsecure string) *exec.Cmd {
  85                 striggerCmdLine = striggerFunc(jobid, containerUUID, finishCommand,
  86                         apiHost, apiToken, apiInsecure).Args
  87                 go func() {
  88                         time.Sleep(5 * time.Second)
  89                         arv.Update("containers", containerUUID,
  90                                 arvadosclient.Dict{
  91                                         "container": arvadosclient.Dict{"state": "Complete"}},
  92                                 nil)
  93                 }()
  94                 return exec.Command("echo", "strigger")
  95         }
  96
  97         go func() {
  98                 time.Sleep(8 * time.Second)
  99                 sigChan <- syscall.SIGINT
 100         }()
 101
 102         // There should be no queued containers now
 103         params := arvadosclient.Dict{
 104                 "filters": [][]string{[]string{"state", "=", "Queued"}},
 105         }
 106         var containers ContainerList
 107         err := arv.List("containers", params, &containers)
 108         c.Check(err, IsNil)
 109         c.Check(len(containers.Items), Equals, 1)
 110
 111         err = doMain()
 112         c.Check(err, IsNil)
 113
 114         c.Check(sbatchCmdLine, DeepEquals, []string{"sbatch", "--job-name=zzzzz-dz642-queuedcontainer", "--share", "--parsable"})
 115         c.Check(striggerCmdLine, DeepEquals, []string{"strigger", "--set", "--jobid=zzzzz-dz642-queuedcontainer\n", "--fini",
 116                 "--program=/usr/bin/crunch-finish-slurm.sh " + os.Getenv("ARVADOS_API_HOST") + " 4axaw8zxe0qm22wa6urpp5nskcne8z88cvbupv653y1njyi05h 1 zzzzz-dz642-queuedcontainer"})
 117
 118         // There should be no queued containers now
 119         err = arv.List("containers", params, &containers)
 120         c.Check(err, IsNil)
 121         c.Check(len(containers.Items), Equals, 0)
 122
 123         // Previously "Queued" container should now be in "Complete" state
 124         var container Container
 125         err = arv.Get("containers", "zzzzz-dz642-queuedcontainer", nil, &container)
 126         c.Check(err, IsNil)
 127         c.Check(container.State, Equals, "Complete")
 128 }
 129
 130 func (s *MockArvadosServerSuite) Test_APIErrorGettingContainers(c *C) {
 131         apiStubResponses := make(map[string]arvadostest.StubResponse)
 132         apiStubResponses["/arvados/v1/containers"] = arvadostest.StubResponse{500, string(`{}`)}
 133
 134         testWithServerStub(c, apiStubResponses, "echo", "Error getting list of queued containers")
 135 }
 136
 137 func testWithServerStub(c *C, apiStubResponses map[string]arvadostest.StubResponse, crunchCmd string, expected string) {
 138         apiStub := arvadostest.ServerStub{apiStubResponses}
 139
 140         api := httptest.NewServer(&apiStub)
 141         defer api.Close()
 142
 143         arv = arvadosclient.ArvadosClient{
 144                 Scheme:    "http",
 145                 ApiServer: api.URL[7:],
 146                 ApiToken:  "abc123",
 147                 Client:    &http.Client{Transport: &http.Transport{}},
 148                 Retries:   0,
 149         }
 150
 151         tempfile, err := ioutil.TempFile(os.TempDir(), "temp-log-file")
 152         c.Check(err, IsNil)
 153         defer os.Remove(tempfile.Name())
 154         log.SetOutput(tempfile)
 155
 156         go func() {
 157                 time.Sleep(2 * time.Second)
 158                 sigChan <- syscall.SIGTERM
 159         }()
 160
 161         runQueuedContainers(2, 1, crunchCmd, crunchCmd)
 162
 163         buf, _ := ioutil.ReadFile(tempfile.Name())
 164         c.Check(strings.Contains(string(buf), expected), Equals, true)
 165 }