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