Fix flaky test

GoogleCloudPlatform · Feb 9, 2016 · a1e882e · a1e882e
1 parent 5de140c
commit a1e882e
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 13 deletions.
diff --git a/managed_vms/sparkjava/src/main/java/com/google/appengine/sparkdemo/Main.java b/managed_vms/sparkjava/src/main/java/com/google/appengine/sparkdemo/Main.java
@@ -27,7 +27,15 @@ public class Main {
    */
   public static void main(String[] args) {
     port(8080);
+    String kind = "DemoUser";
+    if (args != null) {
+      for (String arg : args) {
+        if (arg.startsWith("kind=")) {
+          kind = arg.substring("kind=".length());
+        }
+      }
+    }
     UserController userController =
-        new UserController(new UserService(DatastoreOptions.defaultInstance().service()));
+        new UserController(new UserService(DatastoreOptions.defaultInstance().service(), kind));
   }
 }
diff --git a/managed_vms/sparkjava/src/main/java/com/google/appengine/sparkdemo/UserService.java b/managed_vms/sparkjava/src/main/java/com/google/appengine/sparkdemo/UserService.java
@@ -20,7 +20,6 @@
 
 import com.google.gcloud.datastore.Datastore;
 import com.google.gcloud.datastore.Entity;
-import com.google.gcloud.datastore.FullEntity;
 import com.google.gcloud.datastore.Key;
 import com.google.gcloud.datastore.KeyFactory;
 import com.google.gcloud.datastore.Query;
@@ -31,19 +30,22 @@
 
 public class UserService {
 
-  private static final String KINDNAME = "DEMO_USER";
-  private Datastore datastore;
+  private final Datastore datastore;
+  private final KeyFactory keyFactory;
+  private final String kind;
 
-  public UserService(Datastore datastore) {
+  public UserService(Datastore datastore, String kind) {
     this.datastore = datastore;
+    this.keyFactory = datastore.newKeyFactory().kind(kind);
+    this.kind = kind;
   }
 
   /**
    * Return a list of all users.
    */
   public List<User> getAllUsers() {
     Query<Entity> query =
-        Query.gqlQueryBuilder(Query.ResultType.ENTITY, "SELECT * FROM " + KINDNAME).build();
+        Query.gqlQueryBuilder(Query.ResultType.ENTITY, "SELECT * FROM " + kind).build();
     QueryResults<Entity> results = datastore.run(query);
     List<User> users = new ArrayList<>();
     while (results.hasNext()) {
@@ -60,9 +62,8 @@ public List<User> getAllUsers() {
   public User createUser(String name, String email) {
     failIfInvalid(name, email);
     User user = new User(name, email);
-    KeyFactory keyFactory = datastore.newKeyFactory().kind(KINDNAME);
     Key key = keyFactory.newKey(user.getId());
-    FullEntity entity = Entity.builder(key)
+    Entity entity = Entity.builder(key)
         .set("id", user.getId())
         .set("name", name)
         .set("email", email)
@@ -75,7 +76,6 @@ public User createUser(String name, String email) {
    * Delete a user from Cloud Datastore.
    */
   public String deleteUser(String id) {
-    KeyFactory keyFactory = datastore.newKeyFactory().kind(KINDNAME);
     Key key = keyFactory.newKey(id);
     datastore.delete(key);
     return "ok";
@@ -86,7 +86,6 @@ public String deleteUser(String id) {
    */
   public User updateUser(String id, String name, String email) {
     failIfInvalid(name, email);
-    KeyFactory keyFactory = datastore.newKeyFactory().kind(KINDNAME);
     Key key = keyFactory.newKey(id);
     Entity entity = datastore.get(key);
     if (entity == null) {

diff --git a/managed_vms/sparkjava/src/test/java/com/google/appengine/sparkdemo/UserControllerTest.java b/managed_vms/sparkjava/src/test/java/com/google/appengine/sparkdemo/UserControllerTest.java
@@ -34,6 +34,7 @@
 import java.io.IOException;
 import java.net.HttpURLConnection;
 import java.net.URL;
+import java.util.UUID;
 
 public class UserControllerTest {
 
@@ -43,7 +44,7 @@ public class UserControllerTest {
 
   @BeforeClass
   public static void beforeClass() {
-    Main.main(null);
+    Main.main(new String[] {"kind=DemoUser" + UUID.randomUUID().toString().replaceAll("-", "")});
     Spark.awaitInitialization();
   }
 

diff --git a/managed_vms/sparkjava/src/test/java/com/google/appengine/sparkdemo/UserServiceTest.java b/managed_vms/sparkjava/src/test/java/com/google/appengine/sparkdemo/UserServiceTest.java
@@ -47,7 +47,8 @@ public class UserServiceTest {
   private static final String USER_NAME = "myName";
   private static final String USER_EMAIL = "my@email.com";
   private static final User USER = new User(USER_ID, USER_NAME, USER_EMAIL);
-  private static final Key USER_KEY = Key.builder(PROJECT_ID, "DEMO_USER", USER_ID).build();
+  private static final String KIND = "DemoUser";
+  private static final Key USER_KEY = Key.builder(PROJECT_ID, KIND, USER_ID).build();
   private static final Entity USER_RECORD = Entity.builder(USER_KEY)
       .set("id", USER_ID)
       .set("name", USER_NAME)
@@ -67,7 +68,7 @@ public static void beforeClass() throws IOException, InterruptedException {
         .host("http://localhost:" + PORT)
         .build()
         .service();
-    userService = new UserService(datastore);
+    userService = new UserService(datastore, KIND);
   }
 
   @Before