Add reset() method to LocalDatastoreHelper

mziccard · mziccard · commit 1809167db1df · 2016-09-29T09:23:08.000+02:00
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
@@ -23,6 +23,7 @@
 import com.google.cloud.RetryParams;
 import com.google.cloud.datastore.DatastoreOptions;
 import com.google.common.base.Strings;
+import com.google.common.io.CharStreams;
 
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
@@ -508,6 +509,24 @@ public static boolean sendQuitRequest(int port) {
     return result.toString().startsWith(shutdownMsg);
   }
 
+  public String sendPostRequest(String request) throws IOException {
+    URL url = new URL("http", "localhost", this.port, request);
+    HttpURLConnection con = (HttpURLConnection) url.openConnection();
+    con.setRequestMethod("POST");
+    con.setDoOutput(true);
+    OutputStream out = con.getOutputStream();
+    out.write("".getBytes());
+    out.flush();
+
+    InputStream in = con.getInputStream();
+    String response = CharStreams.toString(new InputStreamReader(con.getInputStream()));
+    in.close();
+    return response;
+  }
+
+  /**
+   * Quit the local emulator and related local service.
+   */
   public void stop() throws IOException, InterruptedException {
     sendQuitRequest(port);
     if (processReader != null) {
@@ -521,6 +540,13 @@ public void stop() throws IOException, InterruptedException {
     }
   }
 
+  /**
+   * Reset the internal state of the emulator.
+   */
+  public void reset() throws IOException {
+    sendPostRequest("/reset");
+  }
+
   private static void deleteRecurse(Path path) throws IOException {
     if (path == null || !Files.exists(path)) {
       return;
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/LocalDatastoreHelperTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/LocalDatastoreHelperTest.java
@@ -17,23 +17,36 @@
 package com.google.cloud.datastore.testing;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.AuthCredentials;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.Key;
 
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.io.IOException;
+
 @RunWith(JUnit4.class)
 public class LocalDatastoreHelperTest {
 
   private static final double TOLERANCE = 0.00001;
   private static final String PROJECT_ID_PREFIX = "test-project-";
   private static final String NAMESPACE = "namespace";
 
+  @Rule
+  public ExpectedException thrown = ExpectedException.none();
+
   @Test
   public void testCreate() {
     LocalDatastoreHelper helper = LocalDatastoreHelper.create(0.75);
@@ -57,4 +70,19 @@ public void testOptions() {
     assertSame(AuthCredentials.noAuth(), options.authCredentials());
     assertEquals(NAMESPACE, options.namespace());
   }
+
+  @Test
+  public void testStartStopReset() throws IOException, InterruptedException {
+    LocalDatastoreHelper helper = LocalDatastoreHelper.create();
+    helper.start();
+    Datastore datastore = helper.options().service();
+    Key key = datastore.newKeyFactory().kind("kind").newKey("name");
+    datastore.put(Entity.builder(key).build());
+    assertNotNull(datastore.get(key));
+    helper.reset();
+    assertNull(datastore.get(key));
+    helper.stop();
+    thrown.expect(DatastoreException.class);
+    datastore.get(key);
+  }
 }
