feat: add support to customize gcloud command of LocalDatastoreHelper

google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java

@@ -35,7 +35,6 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.TimeoutException;
-import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.threeten.bp.Duration;
 
@@ -50,6 +49,7 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
   private final List<EmulatorRunner> emulatorRunners;
   private final double consistency;
   private final Path gcdPath;
+  private boolean storeOnDisk;
 
   // Gcloud emulator settings
   private static final String GCLOUD_CMD_TEXT = "gcloud beta emulators datastore start";
@@ -78,27 +78,65 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
     }
   }
 
-  private LocalDatastoreHelper(double consistency, int port) {
+  /** A builder for {@code LocalDatastoreHelper} objects. */
+  public static class Builder {
+    private double consistency;
+    private int port;
+    private Path dataDir;
+    private boolean storeOnDisk = true;
+
+    private Builder() {}
+
+    private Builder(LocalDatastoreHelper helper) {
+      this.consistency = helper.consistency;
+      this.dataDir = helper.gcdPath;
+      this.storeOnDisk = helper.storeOnDisk;
+    }
+
+    public Builder setConsistency(double consistency) {
+      this.consistency = consistency;
+      return this;
+    }
+
+    public Builder setPort(int port) {
+      this.port = port;
+      return this;
+    }
+
+    public Builder setDataDir(Path dataDir) {
+      this.dataDir = dataDir;
+      return this;
+    }
+
+    public Builder setStoreOnDisk(boolean storeOnDisk) {
+      this.storeOnDisk = storeOnDisk;
+      return this;
+    }
+
+    /** Creates a {@code LocalDatastoreHelper} object. */
+    public LocalDatastoreHelper build() {
+      return new LocalDatastoreHelper(this);
+    }
+  }
+
+  private LocalDatastoreHelper(Builder builder) {
     super(
         "datastore",
-        port > 0 ? port : BaseEmulatorHelper.findAvailablePort(DEFAULT_PORT),
+        builder.port > 0 ? builder.port : BaseEmulatorHelper.findAvailablePort(DEFAULT_PORT),
         PROJECT_ID_PREFIX + UUID.randomUUID().toString());
-    Path tmpDirectory = null;
-    try {
-      tmpDirectory = Files.createTempDirectory("gcd");
-    } catch (IOException ex) {
-      getLogger().log(Level.WARNING, "Failed to create temporary directory");
-    }
-    this.gcdPath = tmpDirectory;
-    this.consistency = consistency;
+    this.consistency = builder.consistency > 0 ? builder.consistency : DEFAULT_CONSISTENCY;
+    this.gcdPath = builder.dataDir;
+    this.storeOnDisk = builder.storeOnDisk;
     String binName = BIN_NAME;
     if (isWindows()) {
       binName = BIN_NAME.replace("/", "\\");
     }
     List<String> gcloudCommand = new ArrayList<>(Arrays.asList(GCLOUD_CMD_TEXT.split(" ")));
     gcloudCommand.add(GCLOUD_CMD_PORT_FLAG + "localhost:" + getPort());
-    gcloudCommand.add(CONSISTENCY_FLAG + consistency);
-    gcloudCommand.add("--no-store-on-disk");
+    gcloudCommand.add(CONSISTENCY_FLAG + builder.consistency);
+    if (!storeOnDisk) {
+      gcloudCommand.add("--no-store-on-disk");
+    }
     GcloudEmulatorRunner gcloudRunner =
         new GcloudEmulatorRunner(gcloudCommand, VERSION_PREFIX, MIN_VERSION);
     List<String> binCommand = new ArrayList<>(Arrays.asList(binName, "start"));
@@ -110,7 +148,17 @@ private LocalDatastoreHelper(double consistency, int port) {
     }
     DownloadableEmulatorRunner downloadRunner =
         new DownloadableEmulatorRunner(binCommand, EMULATOR_URL, MD5_CHECKSUM);
-    emulatorRunners = ImmutableList.of(gcloudRunner, downloadRunner);
+    this.emulatorRunners = ImmutableList.of(gcloudRunner, downloadRunner);
+  }
+
+  /** Returns a builder for {@code LocalDatastoreHelper} object. */
+  public LocalDatastoreHelper.Builder toBuilder() {
+    return new Builder(this);
+  }
+
+  /** Returns a builder for {@code LocalDatastoreHelper} object. */
+  public static LocalDatastoreHelper.Builder newBuilder() {
+    return new LocalDatastoreHelper.Builder();
   }
 
   @Override
@@ -153,6 +201,16 @@ public double getConsistency() {
     return consistency;
   }
 
+  /** Returns the data directory path of the local Datastore emulator. */
+  public Path getGcdPath() {
+    return gcdPath;
+  }
+
+  /** Returns {@code true} data persist on disk, otherwise {@code false} data not store on disk. */
+  public boolean isStoreOnDisk() {
+    return storeOnDisk;
+  }
+
   /**
    * Creates a local Datastore helper with the specified settings for project ID and consistency.
    *
@@ -162,7 +220,7 @@ public double getConsistency() {
    *     consistency of non-ancestor queries; non-ancestor queries are eventually consistent.
    */
   public static LocalDatastoreHelper create(double consistency) {
-    return create(consistency, 0);
+    return LocalDatastoreHelper.newBuilder().setConsistency(consistency).build();
   }
 
   /**
@@ -176,7 +234,7 @@ public static LocalDatastoreHelper create(double consistency) {
    *     emulator will search for a free random port.
    */
   public static LocalDatastoreHelper create(double consistency, int port) {
-    return new LocalDatastoreHelper(consistency, port);
+    return LocalDatastoreHelper.newBuilder().setConsistency(consistency).setPort(port).build();
   }
 
   /**
@@ -187,7 +245,10 @@ public static LocalDatastoreHelper create(double consistency, int port) {
    *     emulator will search for a free random port.
    */
   public static LocalDatastoreHelper create(int port) {
-    return new LocalDatastoreHelper(DEFAULT_CONSISTENCY, port);
+    return LocalDatastoreHelper.newBuilder()
+        .setConsistency(DEFAULT_CONSISTENCY)
+        .setPort(port)
+        .build();
   }
 
   /**
@@ -197,7 +258,7 @@ public static LocalDatastoreHelper create(int port) {
    * all writes are immediately visible.
    */
   public static LocalDatastoreHelper create() {
-    return create(DEFAULT_CONSISTENCY);
+    return LocalDatastoreHelper.newBuilder().build();
   }
 
   /**

google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/ITLocalDatastoreHelperTest.java

@@ -30,6 +30,8 @@
 import com.google.cloud.datastore.Entity;
 import com.google.cloud.datastore.Key;
 import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.concurrent.TimeoutException;
 import org.junit.Assert;
 import org.junit.Test;
@@ -43,6 +45,7 @@ public class ITLocalDatastoreHelperTest {
   private static final double TOLERANCE = 0.00001;
   private static final String PROJECT_ID_PREFIX = "test-project-";
   private static final String NAMESPACE = "namespace";
+  private static final Path DATA_DIR = Paths.get("DATA-DIR");
 
   @Test
   public void testCreate() {
@@ -54,6 +57,56 @@ public void testCreate() {
     assertTrue(helper.getProjectId().startsWith(PROJECT_ID_PREFIX));
   }
 
+  @Test
+  public void testCreateWithBuilder() {
+    LocalDatastoreHelper helper =
+        LocalDatastoreHelper.newBuilder()
+            .setConsistency(0.75)
+            .setPort(8081)
+            .setStoreOnDisk(false)
+            .setDataDir(DATA_DIR)
+            .build();
+    assertTrue(Math.abs(0.75 - helper.getConsistency()) < TOLERANCE);
+    assertTrue(helper.getProjectId().startsWith(PROJECT_ID_PREFIX));
+    assertFalse(helper.isStoreOnDisk());
+    assertEquals(8081, helper.getPort());
+    assertEquals(DATA_DIR, helper.getGcdPath());
+    LocalDatastoreHelper incompleteHelper = LocalDatastoreHelper.newBuilder().build();
+    assertTrue(Math.abs(0.9 - incompleteHelper.getConsistency()) < TOLERANCE);
+    assertTrue(incompleteHelper.getProjectId().startsWith(PROJECT_ID_PREFIX));
+  }
+
+  @Test
+  public void testCreateWithToBuilder() {
+    LocalDatastoreHelper helper =
+        LocalDatastoreHelper.newBuilder()
+            .setConsistency(0.75)
+            .setPort(8081)
+            .setStoreOnDisk(false)
+            .setDataDir(DATA_DIR)
+            .build();
+    assertTrue(Math.abs(0.75 - helper.getConsistency()) < TOLERANCE);
+    assertTrue(helper.getProjectId().startsWith(PROJECT_ID_PREFIX));
+    assertFalse(helper.isStoreOnDisk());
+    assertEquals(8081, helper.getPort());
+    assertEquals(DATA_DIR, helper.getGcdPath());
+    LocalDatastoreHelper actualHelper = helper.toBuilder().build();
+    compareLocalDatastoreHelper(helper, actualHelper);
+    Path dataDir = Paths.get("data-dir");
+    actualHelper =
+        helper
+            .toBuilder()
+            .setConsistency(0.85)
+            .setPort(9091)
+            .setStoreOnDisk(true)
+            .setDataDir(dataDir)
+            .build();
+    assertTrue(Math.abs(0.85 - actualHelper.getConsistency()) < TOLERANCE);
+    assertTrue(actualHelper.isStoreOnDisk());
+    assertEquals(9091, actualHelper.getPort());
+    assertEquals(dataDir, actualHelper.getGcdPath());
+  }
+
   @Test
   public void testCreatePort() {
     LocalDatastoreHelper helper = LocalDatastoreHelper.create(0.75, 8888);
@@ -103,4 +156,31 @@ public void testStartStopReset() throws IOException, InterruptedException, Timeo
       assertNotNull(ex.getMessage());
     }
   }
+
+  @Test
+  public void testStartStopResetWithBuilder()
+      throws IOException, InterruptedException, TimeoutException {
+    try {
+      LocalDatastoreHelper helper = LocalDatastoreHelper.newBuilder().build();
+      helper.start();
+      Datastore datastore = helper.getOptions().getService();
+      Key key = datastore.newKeyFactory().setKind("kind").newKey("name");
+      datastore.put(Entity.newBuilder(key).build());
+      assertNotNull(datastore.get(key));
+      helper.reset();
+      assertNull(datastore.get(key));
+      helper.stop(Duration.ofMinutes(1));
+      datastore.get(key);
+      Assert.fail();
+    } catch (DatastoreException ex) {
+      assertNotNull(ex.getMessage());
+    }
+  }
+
+  public void compareLocalDatastoreHelper(
+      LocalDatastoreHelper expected, LocalDatastoreHelper actual) {
+    assertEquals(expected.getConsistency(), actual.getConsistency(), 0);
+    assertEquals(expected.isStoreOnDisk(), actual.isStoreOnDisk());
+    assertEquals(expected.getGcdPath(), actual.getGcdPath());
+  }
 }