X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a758d03b74f1c62f7ed442f4a4e571e5bba29ddd..a27348943c6e35eee983d6a333eb6977394c3f13:/lib/dispatchcloud/dispatcher_test.go

diff --git a/lib/dispatchcloud/dispatcher_test.go b/lib/dispatchcloud/dispatcher_test.go
index 42767a159a..1f94577434 100644
--- 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()