14324: Use logrus in Azure driver. Fix Sirupsen->sirupsen in imports

[arvados.git] / lib / dispatchcloud / dispatcher_test.go

diff --git a/lib/dispatchcloud/dispatcher_test.go b/lib/dispatchcloud/dispatcher_test.go
index 7d3ab79309567bb853d07daed11b6820dc3cdfca..1f94577434c4076fa0eca25cb2432d579631c657 100644 (file)
--- a/lib/dispatchcloud/dispatcher_test.go
+++ b/lib/dispatchcloud/dispatcher_test.go
@@ -16,7 +16,7 @@ import (
 
        "git.curoverse.com/arvados.git/lib/dispatchcloud/test"
        "git.curoverse.com/arvados.git/sdk/go/arvados"
-       "github.com/Sirupsen/logrus"
+       "github.com/sirupsen/logrus"
        "golang.org/x/crypto/ssh"
        check "gopkg.in/check.v1"
 )
@@ -121,16 +121,18 @@ func (s *DispatcherSuite) TestDispatchToStubDriver(c *check.C) {
        for _, ctr := range queue.Containers {
                waiting[ctr.UUID] = struct{}{}
        }
-       onComplete := func(uuid string) {
+       executeContainer := func(ctr arvados.Container) int {
                mtx.Lock()
                defer mtx.Unlock()
-               if _, ok := waiting[uuid]; !ok {
-                       c.Logf("container completed twice: %s -- perhaps completed after stub instance was killed?", uuid)
+               if _, ok := waiting[ctr.UUID]; !ok {
+                       c.Logf("container completed twice: %s -- perhaps completed after stub instance was killed?", ctr.UUID)
+                       return 1
                }
-               delete(waiting, uuid)
+               delete(waiting, ctr.UUID)
                if len(waiting) == 0 {
                        close(done)
                }
+               return int(rand.Uint32() & 0x3)
        }
        n := 0
        s.stubDriver.Queue = queue
@@ -138,7 +140,7 @@ func (s *DispatcherSuite) TestDispatchToStubDriver(c *check.C) {
                n++
                stubvm.Boot = time.Now().Add(time.Duration(rand.Int63n(int64(5 * time.Millisecond))))
                stubvm.CrunchRunDetachDelay = time.Duration(rand.Int63n(int64(10 * time.Millisecond)))
-               stubvm.CtrExit = int(rand.Uint32() & 0x3)
+               stubvm.ExecuteContainer = executeContainer
                switch n % 7 {
                case 0:
                        stubvm.Broken = time.Now().Add(time.Duration(rand.Int63n(90)) * time.Millisecond)
@@ -147,8 +149,6 @@ func (s *DispatcherSuite) TestDispatchToStubDriver(c *check.C) {
                default:
                        stubvm.CrunchRunCrashRate = 0.1
                }
-               stubvm.OnComplete = onComplete
-               stubvm.OnCancel = onComplete
        }
 
        start := time.Now()
@@ -179,8 +179,12 @@ func (s *DispatcherSuite) TestDispatchToStubDriver(c *check.C) {
 }
 
 func (s *DispatcherSuite) TestAPIPermissions(c *check.C) {
-       drivers["test"] = s.stubDriver
        s.cluster.ManagementToken = "abcdefgh"
+       drivers["test"] = s.stubDriver
+       s.disp.setupOnce.Do(s.disp.initialize)
+       s.disp.queue = &test.Queue{}
+       go s.disp.run()
+
        for _, token := range []string{"abc", ""} {
                req := httptest.NewRequest("GET", "/arvados/v1/dispatch/instances", nil)
                if token != "" {
@@ -197,8 +201,12 @@ func (s *DispatcherSuite) TestAPIPermissions(c *check.C) {
 }
 
 func (s *DispatcherSuite) TestAPIDisabled(c *check.C) {
-       drivers["test"] = s.stubDriver
        s.cluster.ManagementToken = ""
+       drivers["test"] = s.stubDriver
+       s.disp.setupOnce.Do(s.disp.initialize)
+       s.disp.queue = &test.Queue{}
+       go s.disp.run()
+
        for _, token := range []string{"abc", ""} {
                req := httptest.NewRequest("GET", "/arvados/v1/dispatch/instances", nil)
                if token != "" {
@@ -214,6 +222,9 @@ func (s *DispatcherSuite) TestInstancesAPI(c *check.C) {
        s.cluster.ManagementToken = "abcdefgh"
        s.cluster.CloudVMs.TimeoutBooting = arvados.Duration(time.Second)
        drivers["test"] = s.stubDriver
+       s.disp.setupOnce.Do(s.disp.initialize)
+       s.disp.queue = &test.Queue{}
+       go s.disp.run()
 
        type instance struct {
                Instance             string