X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/8da1cf8c69337be51a7ace6923d8c0c0bc7d36e1..00ff9cf6b638f78df82fb5e15ca6788c8d095f40:/lib/dispatchcloud/test/stub_driver.go

diff --git a/lib/dispatchcloud/test/stub_driver.go b/lib/dispatchcloud/test/stub_driver.go
index 9bc1e4c8f7..6e0b129487 100644
--- a/lib/dispatchcloud/test/stub_driver.go
+++ b/lib/dispatchcloud/test/stub_driver.go
@@ -55,6 +55,8 @@ type StubDriver struct {
 	MinTimeBetweenCreateCalls    time.Duration
 	MinTimeBetweenInstancesCalls time.Duration
 
+	QuotaMaxInstances int
+
 	// If true, Create and Destroy calls block until Release() is
 	// called.
 	HoldCloudOps bool
@@ -125,6 +127,9 @@ func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID,
 	if math_rand.Float64() < sis.driver.ErrorRateCreate {
 		return nil, fmt.Errorf("StubInstanceSet: rand < ErrorRateCreate %f", sis.driver.ErrorRateCreate)
 	}
+	if max := sis.driver.QuotaMaxInstances; max > 0 && len(sis.servers) >= max {
+		return nil, QuotaError{fmt.Errorf("StubInstanceSet: reached QuotaMaxInstances %d", max)}
+	}
 	sis.allowCreateCall = time.Now().Add(sis.driver.MinTimeBetweenCreateCalls)
 	ak := sis.driver.AuthorizedKeys
 	if authKey != nil {
@@ -490,3 +495,9 @@ func copyTags(src cloud.InstanceTags) cloud.InstanceTags {
 func (si stubInstance) PriceHistory(arvados.InstanceType) []cloud.InstancePrice {
 	return nil
 }
+
+type QuotaError struct {
+	error
+}
+
+func (QuotaError) IsQuotaError() bool { return true }