X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/e5394906b154b630699c0edd4add36eca34611b3..ac312e0acae4fd114114081c9f4791d05e640831:/lib/dispatchcloud/test/stub_driver.go

diff --git a/lib/dispatchcloud/test/stub_driver.go b/lib/dispatchcloud/test/stub_driver.go
index 826e5c1af3..0a74d97606 100644
--- a/lib/dispatchcloud/test/stub_driver.go
+++ b/lib/dispatchcloud/test/stub_driver.go
@@ -20,6 +20,7 @@ import (
 	"git.arvados.org/arvados.git/lib/cloud"
 	"git.arvados.org/arvados.git/lib/crunchrun"
 	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
 )
@@ -33,7 +34,10 @@ type StubDriver struct {
 	// SetupVM, if set, is called upon creation of each new
 	// StubVM. This is the caller's opportunity to customize the
 	// VM's error rate and other behaviors.
-	SetupVM func(*StubVM)
+	//
+	// If SetupVM returns an error, that error will be returned to
+	// the caller of Create(), and the new VM will be discarded.
+	SetupVM func(*StubVM) error
 
 	// Bugf, if set, is called if a bug is detected in the caller
 	// or stub. Typically set to (*check.C)Errorf. If unset,
@@ -65,7 +69,7 @@ type StubDriver struct {
 }
 
 // InstanceSet returns a new *StubInstanceSet.
-func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, _ cloud.SharedResourceTags, logger logrus.FieldLogger) (cloud.InstanceSet, error) {
+func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, _ cloud.SharedResourceTags, logger logrus.FieldLogger, reg *prometheus.Registry) (cloud.InstanceSet, error) {
 	if sd.holdCloudOps == nil {
 		sd.holdCloudOps = make(chan bool)
 	}
@@ -151,7 +155,10 @@ func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID,
 		Exec:           svm.Exec,
 	}
 	if setup := sis.driver.SetupVM; setup != nil {
-		setup(svm)
+		err := setup(svm)
+		if err != nil {
+			return nil, err
+		}
 	}
 	sis.servers[svm.id] = svm
 	return svm.Instance(), nil
@@ -194,6 +201,12 @@ type RateLimitError struct{ Retry time.Time }
 func (e RateLimitError) Error() string            { return fmt.Sprintf("rate limited until %s", e.Retry) }
 func (e RateLimitError) EarliestRetry() time.Time { return e.Retry }
 
+type CapacityError struct{ InstanceTypeSpecific bool }
+
+func (e CapacityError) Error() string                { return "insufficient capacity" }
+func (e CapacityError) IsCapacityError() bool        { return true }
+func (e CapacityError) IsInstanceTypeSpecific() bool { return e.InstanceTypeSpecific }
+
 // StubVM is a fake server that runs an SSH service. It represents a
 // VM running in a fake cloud.
 //