DatastoreTestRule.java

package xxx;

import com.google.cloud.datastore.Datastore;
import com.google.cloud.datastore.testing.LocalDatastoreHelper;
import com.google.common.flogger.FluentLogger;
import java.io.IOException;
import org.junit.rules.ExternalResource;
import org.threeten.bp.Duration;

/** Datastore helper for unittests. */
public class DatastoreTestRule extends ExternalResource {
  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
  private static final LocalDatastoreHelper datastoreHelper = LocalDatastoreHelper.create();

  private static final int NUM_RETRIES = 3;

  public static Datastore getDatastore() {
    /* When the datastore emulator crashes, we'll get
     * java.net.ConnectException: Connection refused (Connection refused)
     * For now restart the emulator to let the other tests run.
     */
    try {
      datastoreHelper.reset();
    } catch (IOException e) {
      logger.atWarning().withCause(e).log("Failed to reset LocalDatastoreHelper, restarting");
      try {
        datastoreHelper.start();
      } catch (IOException|InterruptedException e2) {
        logger.atSevere().withCause(e2).log("Failed to start LocalDatastoreHelper");
        return null;
      }
    }
    return datastoreHelper.getOptions().getService();
  }

  @Override
  protected void before() throws Throwable {
    datastoreHelper.start();
  }

  @Override
  protected void after() {
    try {
      datastoreHelper.stop(Duration.ofSeconds(10));
    } catch (Exception e) {
      logger.atSevere().withCause(e).log("Failed to stop LocalDatastoreHelper");
    }
  }
}