JCLOUDS-1022: Automatically handle DigitalOcean rate limit

providers/digitalocean2/src/main/java/org/jclouds/digitalocean2/config/DigitalOcean2Properties.java

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.config;
+
+public final class DigitalOcean2Properties {
+
+   /**
+    * Maximum amount of time (in milliseconds) a request will wait until retrying if
+    * the rate limit is exhausted.
+    * <p>
+    * Default value: 2 minutes.
+    */
+   public static final String MAX_RATE_LIMIT_WAIT = "jclouds.max-ratelimit-wait";
+
+   private DigitalOcean2Properties() {
+      throw new AssertionError("intentionally unimplemented");
+   }
+
+}

providers/digitalocean2/src/main/java/org/jclouds/digitalocean2/config/DigitalOcean2RateLimitModule.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.config;
+
+import org.jclouds.digitalocean2.handlers.RateLimitRetryHandler;
+import org.jclouds.http.HttpRetryHandler;
+import org.jclouds.http.annotation.ClientError;
+
+import com.google.inject.AbstractModule;
+
+public class DigitalOcean2RateLimitModule extends AbstractModule {
+   @Override
+   protected void configure() {
+      bind(HttpRetryHandler.class).annotatedWith(ClientError.class).to(RateLimitRetryHandler.class);
+   }
+}

providers/digitalocean2/src/main/java/org/jclouds/digitalocean2/exceptions/DigitalOcean2RateLimitExceededException.java

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.exceptions;
+
+import static org.jclouds.digitalocean2.handlers.RateLimitRetryHandler.millisUntilNextAvailableRequest;
+
+import org.jclouds.http.HttpResponse;
+import org.jclouds.rest.RateLimitExceededException;
+
+import com.google.common.annotations.Beta;
+import com.google.common.base.Predicate;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
+
+/**
+ * Provides detailed information for rate limit exceptions.
+ */
+@Beta
+public class DigitalOcean2RateLimitExceededException extends RateLimitExceededException {
+   private static final long serialVersionUID = 1L;
+   private static final String RATE_LIMIT_HEADER_PREFIX = "RateLimit-";
+
+   private Integer totalRequestsPerHour;
+   private Integer remainingRequests;
+   private Long timeToNextAvailableRequest;
+
+   public DigitalOcean2RateLimitExceededException(HttpResponse response) {
+      super(response.getStatusLine() + "\n" + rateLimitHeaders(response));
+      parseRateLimitInfo(response);
+   }
+
+   public DigitalOcean2RateLimitExceededException(HttpResponse response, Throwable cause) {
+      super(response.getStatusLine() + "\n" + rateLimitHeaders(response), cause);
+      parseRateLimitInfo(response);
+   }
+
+   public Integer totalRequestsPerHour() {
+      return totalRequestsPerHour;
+   }
+
+   public Integer remainingRequests() {
+      return remainingRequests;
+   }
+
+   public Long timeToNextAvailableRequest() {
+      return timeToNextAvailableRequest;
+   }
+
+   private void parseRateLimitInfo(HttpResponse response) {
+      String limit = response.getFirstHeaderOrNull("RateLimit-Limit");
+      String remaining = response.getFirstHeaderOrNull("RateLimit-Remaining");
+      String reset = response.getFirstHeaderOrNull("RateLimit-Reset");
+
+      totalRequestsPerHour = limit == null ? null : Integer.valueOf(limit);
+      remainingRequests = remaining == null ? null : Integer.valueOf(remaining);
+      timeToNextAvailableRequest = reset == null ? null : millisUntilNextAvailableRequest(Long.valueOf(reset));
+   }
+
+   private static Multimap<String, String> rateLimitHeaders(HttpResponse response) {
+      return Multimaps.filterKeys(response.getHeaders(), new Predicate<String>() {
+         @Override
+         public boolean apply(String input) {
+            return input.startsWith(RATE_LIMIT_HEADER_PREFIX);
+         }
+      });
+   }
+}

providers/digitalocean2/src/main/java/org/jclouds/digitalocean2/handlers/DigitalOcean2ErrorHandler.java

@@ -20,6 +20,7 @@ import static org.jclouds.http.HttpUtils.closeClientButKeepContentStream;
 
 import javax.inject.Singleton;
 
+import org.jclouds.digitalocean2.exceptions.DigitalOcean2RateLimitExceededException;
 import org.jclouds.http.HttpCommand;
 import org.jclouds.http.HttpErrorHandler;
 import org.jclouds.http.HttpResponse;
@@ -33,15 +34,16 @@ import org.jclouds.rest.ResourceNotFoundException;
  */
 @Singleton
 public class DigitalOcean2ErrorHandler implements HttpErrorHandler {
+
    public void handleError(HttpCommand command, HttpResponse response) {
       // it is important to always read fully and close streams
       byte[] data = closeClientButKeepContentStream(response);
       String message = data != null ? new String(data) : null;
 
       Exception exception = message != null ? new HttpResponseException(command, response, message)
-              : new HttpResponseException(command, response);
+            : new HttpResponseException(command, response);
       message = message != null ? message : String.format("%s -> %s", command.getCurrentRequest().getRequestLine(),
-              response.getStatusLine());
+            response.getStatusLine());
       switch (response.getStatusCode()) {
          case 400:
             break;
@@ -61,6 +63,9 @@ public class DigitalOcean2ErrorHandler implements HttpErrorHandler {
          case 409:
             exception = new IllegalStateException(message, exception);
             break;
+         case 429:
+            exception = new DigitalOcean2RateLimitExceededException(response, exception);
+            break;
       }
       command.setException(exception);
    }

providers/digitalocean2/src/main/java/org/jclouds/digitalocean2/handlers/RateLimitRetryHandler.java

@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.handlers;
+
+import static org.jclouds.Constants.PROPERTY_MAX_RETRIES;
+import static org.jclouds.digitalocean2.config.DigitalOcean2Properties.MAX_RATE_LIMIT_WAIT;
+
+import javax.annotation.Resource;
+import javax.inject.Named;
+import javax.inject.Singleton;
+
+import org.jclouds.http.HttpCommand;
+import org.jclouds.http.HttpResponse;
+import org.jclouds.http.HttpRetryHandler;
+import org.jclouds.logging.Logger;
+
+import com.google.common.annotations.Beta;
+import com.google.inject.Inject;
+
+/**
+ * Retry handler that takes into account the DigitalOcean rate limit and delays
+ * the requests until they are known to succeed.
+ */
+@Beta
+@Singleton
+public class RateLimitRetryHandler implements HttpRetryHandler {
+
+   static final String RATE_LIMIT_RESET_HEADER = "RateLimit-Reset";
+
+   @Resource
+   protected Logger logger = Logger.NULL;
+
+   @Inject(optional = true)
+   @Named(PROPERTY_MAX_RETRIES)
+   private int retryCountLimit = 5;
+
+   @Inject(optional = true)
+   @Named(MAX_RATE_LIMIT_WAIT)
+   private int maxRateLimitWait = 120000;
+
+   @Override
+   public boolean shouldRetryRequest(final HttpCommand command, final HttpResponse response) {
+      command.incrementFailureCount();
+
+      // Do not retry client errors that are not rate limit errors
+      if (response.getStatusCode() != 429) {
+         return false;
+      } else if (!command.isReplayable()) {
+         logger.error("Cannot retry after rate limit error, command is not replayable: %1$s", command);
+         return false;
+      } else if (command.getFailureCount() > retryCountLimit) {
+         logger.error("Cannot retry after rate limit error, command has exceeded retry limit %1$d: %2$s",
+               retryCountLimit, command);
+         return false;
+      } else {
+         return delayRequestUntilAllowed(command, response);
+      }
+   }
+
+   private boolean delayRequestUntilAllowed(final HttpCommand command, final HttpResponse response) {
+      // The header is the Unix epoch time when the next request can be done
+      String epochForNextAvailableRequest = response.getFirstHeaderOrNull(RATE_LIMIT_RESET_HEADER);
+      if (epochForNextAvailableRequest == null) {
+         logger.error("Cannot retry after rate limit error, no retry information provided in the response");
+         return false;
+      }
+
+      long waitPeriod = millisUntilNextAvailableRequest(Long.parseLong(epochForNextAvailableRequest));
+
+      if (waitPeriod > 0) {
+         if (waitPeriod > maxRateLimitWait) {
+            logger.error("Max wait for rate limited requests is %s seconds but need to wait %s seconds, aborting",
+                  maxRateLimitWait, waitPeriod);
+            return false;
+         }
+
+         try {
+            logger.debug("Waiting %s seconds before retrying, as defined by the rate limit", waitPeriod);
+            // Do not use Uninterrumpibles or similar, to let the jclouds
+            // tiemout configuration interrupt this thread
+            Thread.sleep(waitPeriod);
+         } catch (InterruptedException ex) {
+            // If the request is being executed and has a timeout configured,
+            // the thread may be interrupted when the timeout is reached.
+            logger.error("Request execution was interrupted, aborting");
+            Thread.currentThread().interrupt();
+            return false;
+         }
+      }
+
+      return true;
+   }
+
+   public static long millisUntilNextAvailableRequest(long epochForNextAvailableRequest) {
+      return (epochForNextAvailableRequest * 1000) - System.currentTimeMillis();
+   }
+}

providers/digitalocean2/src/test/java/org/jclouds/digitalocean2/compute/DigitalOcean2ComputeServiceLiveTest.java

@@ -18,8 +18,10 @@ package org.jclouds.digitalocean2.compute;
 
 import org.jclouds.compute.domain.NodeMetadata;
 import org.jclouds.compute.internal.BaseComputeServiceLiveTest;
+import org.jclouds.digitalocean2.config.DigitalOcean2RateLimitModule;
 import org.jclouds.sshj.config.SshjSshClientModule;
 import org.testng.annotations.Test;
+
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Module;
@@ -39,6 +41,12 @@ public class DigitalOcean2ComputeServiceLiveTest extends BaseComputeServiceLiveT
       return new SshjSshClientModule();
    }
 
+   @Override
+   protected Iterable<Module> setupModules() {
+      return ImmutableSet.<Module> builder().addAll(super.setupModules()).add(new DigitalOcean2RateLimitModule())
+            .build();
+   }
+
    @Override
    public void testOptionToNotBlock() throws Exception {
       // DigitalOcean ComputeService implementation has to block until the node

providers/digitalocean2/src/test/java/org/jclouds/digitalocean2/exceptions/RateLimitExceptionMockTest.java

@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.exceptions;
+
+import static org.jclouds.Constants.PROPERTY_MAX_RETRIES;
+import static org.jclouds.digitalocean2.handlers.RateLimitRetryHandler.millisUntilNextAvailableRequest;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
+import java.util.Properties;
+
+import org.jclouds.digitalocean2.internal.BaseDigitalOcean2ApiMockTest;
+import org.testng.annotations.Test;
+
+import com.squareup.okhttp.mockwebserver.MockResponse;
+
+@Test(groups = "unit", testName = "RateLimitExceptionMockTest", singleThreaded = true)
+public class RateLimitExceptionMockTest extends BaseDigitalOcean2ApiMockTest {
+
+   @Override
+   protected Properties overrides() {
+      Properties overrides = super.overrides();
+      overrides.put(PROPERTY_MAX_RETRIES, "0"); // Do not retry
+      return overrides;
+   }
+
+   public void testRateLimitExceptionIsThrown() throws InterruptedException {
+      long reset = (System.currentTimeMillis() / 1000) + 3600; // Epoch for one
+                                                               // hour from now
+      long millisToReset = millisUntilNextAvailableRequest(reset);
+
+      server.enqueue(new MockResponse().setResponseCode(429).addHeader("RateLimit-Limit", "5000")
+            .addHeader("RateLimit-Remaining", "1235").addHeader("RateLimit-Reset", String.valueOf(reset)));
+
+      try {
+         api.keyApi().list();
+         fail("Expected a DigitalOcean2RateLimitExceededException to be thrown");
+      } catch (DigitalOcean2RateLimitExceededException ex) {
+         assertEquals(ex.totalRequestsPerHour().intValue(), 5000);
+         assertEquals(ex.remainingRequests().intValue(), 1235);
+         // Can't verify with millisecond precision. Use an interval to have a
+         // consistent test.
+         assertTrue(ex.timeToNextAvailableRequest() < millisToReset
+               && ex.timeToNextAvailableRequest() > millisToReset - 1800000);
+      }
+   }
+
+}

providers/digitalocean2/src/test/java/org/jclouds/digitalocean2/handlers/RateLimitRetryHandlerTest.java

@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jclouds.digitalocean2.handlers;
+
+import static org.jclouds.digitalocean2.handlers.RateLimitRetryHandler.RATE_LIMIT_RESET_HEADER;
+import static org.jclouds.http.HttpUtils.releasePayload;
+import static org.jclouds.io.Payloads.newInputStreamPayload;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+import java.io.ByteArrayInputStream;
+import java.util.concurrent.TimeUnit;
+
+import org.jclouds.http.HttpCommand;
+import org.jclouds.http.HttpRequest;
+import org.jclouds.http.HttpResponse;
+import org.jclouds.io.Payload;
+import org.testng.annotations.Test;
+
+import com.google.common.util.concurrent.Uninterruptibles;
+
+@Test(groups = "unit", testName = "RateLimitRetryHandlerTest")
+public class RateLimitRetryHandlerTest {
+
+   // Configure a safe timeout of one minute to abort the tests in case they get
+   // stuck
+   private static final long TEST_SAFE_TIMEOUT = 60000;
+
+   private final RateLimitRetryHandler rateLimitRetryHandler = new RateLimitRetryHandler();
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDoNotRetryIfNoRateLimit() {
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(450).build();
+
+      assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDoNotRetryIfNotReplayable() {
+      // InputStream payloads are not replayable
+      Payload payload = newInputStreamPayload(new ByteArrayInputStream(new byte[0]));
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost")
+            .payload(payload).build());
+      HttpResponse response = HttpResponse.builder().statusCode(429).build();
+
+      try {
+         assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+      } finally {
+         releasePayload(command.getCurrentRequest());
+      }
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDoNotRetryIfNoRateLimitResetHeader() {
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(429).build();
+
+      assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDoNotRetryIfTooMuchWait() {
+      // 5 minutes Unix epoch timestamp
+      long rateLimitResetEpoch = (System.currentTimeMillis() + 300000) / 1000;
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(429)
+            .addHeader(RATE_LIMIT_RESET_HEADER, String.valueOf(rateLimitResetEpoch)).build();
+
+      assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testRequestIsDelayed() {
+      // 5 seconds Unix epoch timestamp
+      long rateLimitResetEpoch = (System.currentTimeMillis() + 5000) / 1000;
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(429)
+            .addHeader(RATE_LIMIT_RESET_HEADER, String.valueOf(rateLimitResetEpoch)).build();
+
+      long start = System.currentTimeMillis();
+
+      assertTrue(rateLimitRetryHandler.shouldRetryRequest(command, response));
+      // Should have blocked the amount of time configured in the header. Use a
+      // smaller value to compensate the time it takes to reach the code that
+      // computes the amount of time to wait.
+      assertTrue(System.currentTimeMillis() - start > 2500);
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDoNotRetryIfRequestIsAborted() throws Exception {
+      // 10 seconds Unix epoch timestamp
+      long rateLimitResetEpoch = (System.currentTimeMillis() + 10000) / 1000;
+      final HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost")
+            .build());
+      final HttpResponse response = HttpResponse.builder().statusCode(429)
+            .addHeader(RATE_LIMIT_RESET_HEADER, String.valueOf(rateLimitResetEpoch)).build();
+
+      final Thread requestThread = Thread.currentThread();
+      Thread killer = new Thread() {
+         @Override
+         public void run() {
+            Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
+            requestThread.interrupt();
+         }
+      };
+
+      // Start the killer thread that will abort the rate limit wait
+      killer.start();
+      assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testIncrementsFailureCount() {
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(429).build();
+
+      rateLimitRetryHandler.shouldRetryRequest(command, response);
+      assertEquals(command.getFailureCount(), 1);
+
+      rateLimitRetryHandler.shouldRetryRequest(command, response);
+      assertEquals(command.getFailureCount(), 2);
+
+      rateLimitRetryHandler.shouldRetryRequest(command, response);
+      assertEquals(command.getFailureCount(), 3);
+   }
+
+   @Test(timeOut = TEST_SAFE_TIMEOUT)
+   public void testDisallowExcessiveRetries() {
+      HttpCommand command = new HttpCommand(HttpRequest.builder().method("GET").endpoint("http://localhost").build());
+      HttpResponse response = HttpResponse.builder().statusCode(429).addHeader(RATE_LIMIT_RESET_HEADER, "0").build();
+
+      for (int i = 0; i < 5; i++) {
+         assertTrue(rateLimitRetryHandler.shouldRetryRequest(command, response));
+      }
+      assertFalse(rateLimitRetryHandler.shouldRetryRequest(command, response));
+   }
+}

providers/digitalocean2/src/test/java/org/jclouds/digitalocean2/internal/BaseDigitalOcean2ApiLiveTest.java

@@ -28,6 +28,7 @@ import org.jclouds.apis.BaseApiLiveTest;
 import org.jclouds.compute.config.ComputeServiceProperties;
 import org.jclouds.compute.domain.NodeMetadata;
 import org.jclouds.digitalocean2.DigitalOcean2Api;
+import org.jclouds.digitalocean2.config.DigitalOcean2RateLimitModule;
 import org.jclouds.digitalocean2.domain.Action;
 import org.jclouds.digitalocean2.domain.Image;
 import org.jclouds.digitalocean2.domain.Region;
@@ -35,6 +36,7 @@ import org.jclouds.digitalocean2.domain.Size;
 
 import com.google.common.base.Predicate;
 import com.google.common.collect.ComparisonChain;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import com.google.inject.Injector;
 import com.google.inject.Key;
@@ -68,6 +70,11 @@ public class BaseDigitalOcean2ApiLiveTest extends BaseApiLiveTest<DigitalOcean2A
       return injector.getInstance(DigitalOcean2Api.class);
    }
 
+   @Override protected Iterable<Module> setupModules() {
+      return ImmutableSet.<Module> builder().addAll(super.setupModules()).add(new DigitalOcean2RateLimitModule())
+            .build();
+   }
+
    protected void assertActionCompleted(int actionId) {
       checkState(actionCompleted.apply(actionId), "Timeout waiting for action: %s", actionId);
       Action action = api.actionApi().get(actionId);

providers/digitalocean2/src/test/java/org/jclouds/digitalocean2/internal/BaseDigitalOcean2ApiMockTest.java

@@ -23,6 +23,7 @@ import static org.testng.Assert.assertEquals;
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.Properties;
 import java.util.Set;
 
 import org.jclouds.ContextBuilder;
@@ -59,7 +60,6 @@ public class BaseDigitalOcean2ApiMockTest {
    // So that we can ignore formatting.
    private final JsonParser parser = new JsonParser();
    
-   
    @BeforeMethod
    public void start() throws IOException {
       server = new MockWebServer();
@@ -68,6 +68,7 @@ public class BaseDigitalOcean2ApiMockTest {
             .credentials("", MOCK_BEARER_TOKEN)
             .endpoint(url(""))
             .modules(modules)
+            .overrides(overrides())
             .build();
       json = ctx.utils().injector().getInstance(Json.class);
       api = ctx.getApi();
@@ -79,6 +80,10 @@ public class BaseDigitalOcean2ApiMockTest {
       api.close();
    }
    
+   protected Properties overrides() {
+      return new Properties();
+   }
+
    protected String url(String path) {
       return server.getUrl(path).toString();
    }