Repository: firebase/firebase-admin-java
Pull request #100 (2017-11-22 19:40:21)

Improving Unit Test Coverage
This PR adds more unit tests to cover classes like `Connection`, `PersistentConnectionImpl` and `WebsocketConnection`. To enable this, I also made some minor modifications to the classes in question, so they are more amenable to unit testing. Overall this brings our unit test coverage to 80%.

hiranya911: I've made all the changes except for the use of generics for different factory types. I've left a follow up question regarding that.

Patch:
diff --git a/src/main/java/com/google/firebase/cloud/StorageClient.java b/src/main/java/com/google/firebase/cloud/StorageClient.java
index c4c2605..96b53c8 100644
--- a/src/main/java/com/google/firebase/cloud/StorageClient.java
+++ b/src/main/java/com/google/firebase/cloud/StorageClient.java
@@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import com.google.cloud.storage.Bucket;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.StorageOptions;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.ImplFirebaseTrampolines;
@@ -41,12 +42,10 @@ public class StorageClient {
   private final FirebaseApp app;
   private final Storage storage;
 
-  private StorageClient(FirebaseApp app) {
+  @VisibleForTesting
+  StorageClient(FirebaseApp app, Storage storage) {
     this.app = checkNotNull(app, "FirebaseApp must not be null");
-    this.storage = StorageOptions.newBuilder()
-        .setCredentials(ImplFirebaseTrampolines.getCredentials(app))
-        .build()
-        .getService();
+    this.storage = checkNotNull(storage, "Storage must not be null");
   }
 
   public static StorageClient getInstance() {
@@ -57,7 +56,12 @@ public class StorageClient {
     StorageClientService service = ImplFirebaseTrampolines.getService(app, SERVICE_ID,
         StorageClientService.class);
     if (service == null) {
-      service = ImplFirebaseTrampolines.addService(app, new StorageClientService(app));
+      Storage storage = StorageOptions.newBuilder()
+          .setCredentials(ImplFirebaseTrampolines.getCredentials(app))
+          .build()
+          .getService();
+      StorageClient client = new StorageClient(app, storage);
+      service = ImplFirebaseTrampolines.addService(app, new StorageClientService(client));
     }
     return service.getInstance();
   }
@@ -99,8 +103,8 @@ public class StorageClient {
 
   private static class StorageClientService extends FirebaseService<StorageClient> {
 
-    StorageClientService(FirebaseApp app) {
-      super(SERVICE_ID, new StorageClient(app));
+    StorageClientService(StorageClient client) {
+      super(SERVICE_ID, client);
     }
 
     @Override
diff --git a/src/main/java/com/google/firebase/database/connection/Connection.java b/src/main/java/com/google/firebase/database/connection/Connection.java
index a64a162..4059e46 100644
--- a/src/main/java/com/google/firebase/database/connection/Connection.java
+++ b/src/main/java/com/google/firebase/database/connection/Connection.java
@@ -16,6 +16,7 @@
 
 package com.google.firebase.database.connection;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.firebase.database.logging.LogWrapper;
 
 import java.util.HashMap;
@@ -39,24 +40,36 @@ class Connection implements WebsocketConnection.Delegate {
   private static final String SERVER_HELLO_HOST = "h";
   private static final String SERVER_HELLO_SESSION_ID = "s";
   private static long connectionIds = 0;
+
   private final LogWrapper logger;
-  private HostInfo hostInfo;
+  private final HostInfo hostInfo;
+  private final Delegate delegate;
+
   private WebsocketConnection conn;
-  private Delegate delegate;
   private State state;
 
-  public Connection(
+  Connection(
       ConnectionContext context,
       HostInfo hostInfo,
       String cachedHost,
       Delegate delegate,
       String optLastSessionId) {
+    this(context, hostInfo, delegate,
+        new DefaultWebsocketConnectionFactory(context, hostInfo, cachedHost, optLastSessionId));
+  }
+
+  @VisibleForTesting
+  Connection(
+      ConnectionContext context,
+      HostInfo hostInfo,
+      Delegate delegate,
+      WebsocketConnectionFactory connFactory) {
     long connId = connectionIds++;
     this.hostInfo = hostInfo;
     this.delegate = delegate;
     this.logger = new LogWrapper(context.getLogger(), Connection.class, "conn_" + connId);
     this.state = State.REALTIME_CONNECTING;
-    this.conn = new WebsocketConnection(context, hostInfo, cachedHost, this, optLastSessionId);
+    this.conn = connFactory.newConnection(this);
   }
 
   public void open() {
@@ -246,11 +259,6 @@ class Connection implements WebsocketConnection.Delegate {
     }
   }
 
-  // For testing
-  public void injectConnectionFailure() {
-    this.close();
-  }
-
   public enum DisconnectReason {
     SERVER_RESET,
     OTHER
@@ -262,6 +270,34 @@ class Connection implements WebsocketConnection.Delegate {
     REALTIME_DISCONNECTED
   }
 
+  interface WebsocketConnectionFactory {
+    WebsocketConnection newConnection(WebsocketConnection.Delegate delegate);
+  }
+
+  private static class DefaultWebsocketConnectionFactory implements WebsocketConnectionFactory {
+
+    final ConnectionContext context;
+    final HostInfo hostInfo;
+    final String cachedHost;
+    final String optLastSessionId;
+
+    DefaultWebsocketConnectionFactory(
+        ConnectionContext context,
+        HostInfo hostInfo,
+        String cachedHost,
+        String optLastSessionId) {
+      this.context = context;
+      this.hostInfo = hostInfo;
+      this.cachedHost = cachedHost;
+      this.optLastSessionId = optLastSessionId;
+    }
+
+    @Override
+    public WebsocketConnection newConnection(WebsocketConnection.Delegate delegate) {
+      return new WebsocketConnection(context, hostInfo, cachedHost, delegate, optLastSessionId);
+    }
+  }
+
   public interface Delegate {
 
     void onCacheHost(String host);
diff --git a/src/main/java/com/google/firebase/database/connection/PersistentConnectionImpl.java b/src/main/java/com/google/firebase/database/connection/PersistentConnectionImpl.java
index 46e35bc..5f2f713 100644
--- a/src/main/java/com/google/firebase/database/connection/PersistentConnectionImpl.java
+++ b/src/main/java/com/google/firebase/database/connection/PersistentConnectionImpl.java
@@ -20,7 +20,6 @@ import static com.google.firebase.database.connection.ConnectionUtils.hardAssert
 
 import com.google.firebase.database.connection.util.RetryHelper;
 import com.google.firebase.database.logging.LogWrapper;
-import com.google.firebase.database.util.AndroidSupport;
 import com.google.firebase.database.util.GAuthToken;
 
 import java.util.ArrayList;
@@ -93,13 +92,16 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
   private static final String IDLE_INTERRUPT_REASON = "connection_idle";
   private static final String TOKEN_REFRESH_INTERRUPT_REASON = "token_refresh";
   private static long connectionIds = 0;
+
   private final Delegate delegate;
   private final HostInfo hostInfo;
   private final ConnectionContext context;
+  private final ConnectionFactory connFactory;
   private final ConnectionAuthTokenProvider authTokenProvider;
   private final ScheduledExecutorService executorService;
   private final LogWrapper logger;
   private final RetryHelper retryHelper;
+
   private String cachedHost;
   private HashSet<String> interruptReasons = new HashSet<>();
   private boolean firstConnection = true;
@@ -123,13 +125,19 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
   private long lastWriteTimestamp;
   private boolean hasOnDisconnects;
 
-  public PersistentConnectionImpl(
-      ConnectionContext context, HostInfo info, final Delegate delegate) {
-    this.delegate = delegate;
+  public PersistentConnectionImpl(ConnectionContext context, HostInfo info, Delegate delegate) {
+    this(context, info, delegate, new DefaultConnectionFactory());
+  }
+
+  PersistentConnectionImpl(
+      ConnectionContext context, HostInfo info, Delegate delegate, ConnectionFactory connFactory) {
+
     this.context = context;
+    this.hostInfo = info;
+    this.delegate = delegate;
+    this.connFactory = connFactory;
     this.executorService = context.getExecutorService();
     this.authTokenProvider = context.getAuthTokenProvider();
-    this.hostInfo = info;
     this.listens = new HashMap<>();
     this.requestCBHash = new HashMap<>();
     this.outstandingPuts = new HashMap<>();
@@ -422,7 +430,7 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
     return interruptReasons.contains(reason);
   }
 
-  boolean shouldReconnect() {
+  private boolean shouldReconnect() {
     return interruptReasons.size() == 0;
   }
 
@@ -524,20 +532,19 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
     }
   }
 
-  public void openNetworkConnection(String token) {
+  private void openNetworkConnection(String token) {
     hardAssert(
-        this.connectionState == ConnectionState.GettingToken,
+        connectionState == ConnectionState.GettingToken,
         "Trying to open network connection while in the wrong state: %s",
-        this.connectionState);
+        connectionState);
     // User might have logged out. Positive auth status is handled after authenticating with
     // the server
     if (token == null) {
-      this.delegate.onAuthStatus(false);
+      delegate.onAuthStatus(false);
     }
-    this.authToken = token;
-    this.connectionState = ConnectionState.Connecting;
-    realtime =
-        new Connection(this.context, this.hostInfo, this.cachedHost, this, this.lastSessionId);
+    authToken = token;
+    connectionState = ConnectionState.Connecting;
+    realtime = connFactory.newConnection(this);
     realtime.open();
   }
 
@@ -1054,17 +1061,9 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
 
   private void sendConnectStats() {
     Map<String, Integer> stats = new HashMap<>();
-    if (AndroidSupport.isAndroid()) {
-      if (this.context.isPersistenceEnabled()) {
-        stats.put("persistence.android.enabled", 1);
-      }
-      stats.put("sdk.android." + context.getClientSdkVersion().replace('.', '-'), 1);
-      // TODO: Also send stats for connection version
-    } else {
-      assert !this.context.isPersistenceEnabled()
-          : "Stats for persistence on JVM missing (persistence not yet supported)";
-      stats.put("sdk.admin_java." + context.getClientSdkVersion().replace('.', '-'), 1);
-    }
+    assert !this.context.isPersistenceEnabled()
+        : "Stats for persistence on JVM missing (persistence not yet supported)";
+    stats.put("sdk.admin_java." + context.getClientSdkVersion().replace('.', '-'), 1);
     if (logger.logsDebug()) {
       logger.debug("Sending first connection stats");
     }
@@ -1137,13 +1136,6 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
     return isIdle() && now > (this.lastWriteTimestamp + IDLE_TIMEOUT);
   }
 
-  // For testing
-  public void injectConnectionFailure() {
-    if (this.realtime != null) {
-      this.realtime.injectConnectionFailure();
-    }
-  }
-
   private enum ConnectionState {
     Disconnected,
     GettingToken,
@@ -1215,15 +1207,15 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
       this.tag = tag;
     }
 
-    public ListenQuerySpec getQuery() {
+    ListenQuerySpec getQuery() {
       return query;
     }
 
-    public Long getTag() {
+    Long getTag() {
       return this.tag;
     }
 
-    public ListenHashProvider getHashFunction() {
+    ListenHashProvider getHashFunction() {
       return this.hashFunction;
     }
 
@@ -1247,23 +1239,23 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
       this.onComplete = onComplete;
     }
 
-    public String getAction() {
+    String getAction() {
       return action;
     }
 
-    public Map<String, Object> getRequest() {
+    Map<String, Object> getRequest() {
       return request;
     }
 
-    public RequestResultCallback getOnComplete() {
+    RequestResultCallback getOnComplete() {
       return onComplete;
     }
 
-    public void markSent() {
+    void markSent() {
       this.sent = true;
     }
 
-    public boolean wasSent() {
+    boolean wasSent() {
       return this.sent;
     }
   }
@@ -1299,4 +1291,18 @@ public class PersistentConnectionImpl implements Connection.Delegate, Persistent
       return onComplete;
     }
   }
+
+  interface ConnectionFactory {
+    Connection newConnection(PersistentConnectionImpl delegate);
+  }
+
+  private static class DefaultConnectionFactory implements ConnectionFactory {
+    @Override
+    public Connection newConnection(PersistentConnectionImpl delegate) {
+      return new Connection(delegate.context, delegate.hostInfo, delegate.cachedHost,
+          delegate, delegate.lastSessionId);
+    }
+  }
+
+
 }
diff --git a/src/main/java/com/google/firebase/database/connection/WebsocketConnection.java b/src/main/java/com/google/firebase/database/connection/WebsocketConnection.java
index 95a372d..d58a1af 100644
--- a/src/main/java/com/google/firebase/database/connection/WebsocketConnection.java
+++ b/src/main/java/com/google/firebase/database/connection/WebsocketConnection.java
@@ -48,7 +48,6 @@ class WebsocketConnection {
   private static final int MAX_FRAME_SIZE = 16384;
   private static final AtomicLong CONN_ID = new AtomicLong(0);
 
-  private final ConnectionContext connectionContext;
   private final ScheduledExecutorService executorService;
   private final LogWrapper logger;
   private final WSClient conn;
@@ -66,22 +65,19 @@ class WebsocketConnection {
       String optCachedHost,
       Delegate delegate,
       String optLastSessionId) {
-    this.connectionContext = connectionContext;
+    this(connectionContext, delegate,
+        new DefaultWSClientFactory(connectionContext, hostInfo, optCachedHost, optLastSessionId));
+  }
+
+  WebsocketConnection(
+      ConnectionContext connectionContext,
+      Delegate delegate,
+      WSClientFactory clientFactory) {
     this.executorService = connectionContext.getExecutorService();
     this.delegate = delegate;
     this.logger = new LogWrapper(connectionContext.getLogger(), WebsocketConnection.class,
         "ws_" + CONN_ID.getAndIncrement());
-    this.conn = createConnection(hostInfo, optCachedHost, optLastSessionId);
-  }
-
-  private WSClient createConnection(
-      HostInfo hostInfo, String optCachedHost, String optLastSessionId) {
-    String host = (optCachedHost != null) ? optCachedHost : hostInfo.getHost();
-    URI uri = HostInfo.getConnectionUrl(
-        host, hostInfo.isSecure(), hostInfo.getNamespace(), optLastSessionId);
-    return new NettyWebSocketClient(
-        uri, connectionContext.getUserAgent(), connectionContext.getThreadFactory(),
-        new WSClientHandlerImpl());
+    this.conn = clientFactory.newClient(new WSClientHandlerImpl());
   }
 
   void open() {
@@ -409,6 +405,37 @@ class WebsocketConnection {
     void send(String msg);
   }
 
+  interface WSClientFactory {
+
+    WSClient newClient(
+        WSClientEventHandler delegate);
+
+  }
+
+  private static class DefaultWSClientFactory implements WSClientFactory {
+
+    final ConnectionContext context;
+    final HostInfo hostInfo;
+    final String optCachedHost;
+    final String optLastSessionId;
+
+    DefaultWSClientFactory(ConnectionContext context, HostInfo hostInfo, String
+        optCachedHost, String optLastSessionId) {
+      this.context = context;
+      this.hostInfo = hostInfo;
+      this.optCachedHost = optCachedHost;
+      this.optLastSessionId = optLastSessionId;
+    }
+
+    public WSClient newClient(WSClientEventHandler delegate) {
+      String host = (optCachedHost != null) ? optCachedHost : hostInfo.getHost();
+      URI uri = HostInfo.getConnectionUrl(
+          host, hostInfo.isSecure(), hostInfo.getNamespace(), optLastSessionId);
+      return new NettyWebSocketClient(
+          uri, context.getUserAgent(), context.getThreadFactory(), delegate);
+    }
+  }
+
   /**
    * Event handler that handles the events generated by a low-level {@link WSClient}.
    */
diff --git a/src/main/java/com/google/firebase/database/core/Repo.java b/src/main/java/com/google/firebase/database/core/Repo.java
index 8d3bc83..add6ff3 100644
--- a/src/main/java/com/google/firebase/database/core/Repo.java
+++ b/src/main/java/com/google/firebase/database/core/Repo.java
@@ -75,7 +75,6 @@ public class Repo implements PersistentConnection.Delegate {
   private final LogWrapper operationLogger;
   private final LogWrapper transactionLogger;
   private final LogWrapper dataLogger;
-  public long dataUpdateCount = 0; // for testing.
   private SnapshotHolder infoData;
   private SparseSnapshotTree onDisconnect;
   private Tree<List<TransactionData>> transactionQueueTree;
@@ -216,50 +215,7 @@ public class Repo implements PersistentConnection.Delegate {
   }
 
   private void restoreWrites(PersistenceManager persistenceManager) {
-    List<UserWriteRecord> writes = persistenceManager.loadUserWrites();
-
-    Map<String, Object> serverValues = ServerValues.generateServerValues(serverClock);
-    long lastWriteId = Long.MIN_VALUE;
-    for (final UserWriteRecord write : writes) {
-      RequestResultCallback onComplete =
-          new RequestResultCallback() {
-        @Override
-        public void onRequestResult(String optErrorCode, String optErrorMessage) {
-          DatabaseError error = fromErrorCode(optErrorCode, optErrorMessage);
-          warnIfWriteFailed("Persisted write", write.getPath(), error);
-          ackWriteAndRerunTransactions(write.getWriteId(), write.getPath(), error);
-        }
-      };
-      if (lastWriteId >= write.getWriteId()) {
-        throw new IllegalStateException("Write ids were not in order.");
-      }
-      lastWriteId = write.getWriteId();
-      nextWriteId = write.getWriteId() + 1;
-      if (write.isOverwrite()) {
-        if (operationLogger.logsDebug()) {
-          operationLogger.debug("Restoring overwrite with id " + write.getWriteId());
-        }
-        connection.put(write.getPath().asList(), write.getOverwrite().getValue(true), onComplete);
-        Node resolved =
-            ServerValues.resolveDeferredValueSnapshot(write.getOverwrite(), serverValues);
-        serverSyncTree.applyUserOverwrite(
-            write.getPath(),
-            write.getOverwrite(),
-            resolved,
-            write.getWriteId(), /*visible=*/
-            true, /*persist=*/
-            false);
-      } else {
-        if (operationLogger.logsDebug()) {
-          operationLogger.debug("Restoring merge with id " + write.getWriteId());
-        }
-        connection.merge(write.getPath().asList(), write.getMerge().getValue(true), onComplete);
-        CompoundWrite resolved =
-            ServerValues.resolveDeferredValueMerge(write.getMerge(), serverValues);
-        serverSyncTree.applyUserMerge(
-            write.getPath(), write.getMerge(), resolved, write.getWriteId(), /*persist=*/ false);
-      }
-    }
+    // No-op (Admin SDK does not support persistence)
   }
 
   public FirebaseDatabase getDatabase() {
@@ -281,7 +237,7 @@ public class Repo implements PersistentConnection.Delegate {
     ctx.getRunLoop().scheduleNow(r);
   }
 
-  public void postEvent(Runnable r) {
+  private void postEvent(Runnable r) {
     InternalHelpers.checkNotDestroyed(this);
     ctx.requireStarted();
     ctx.getEventTarget().postEvent(r);
@@ -313,7 +269,6 @@ public class Repo implements PersistentConnection.Delegate {
     if (dataLogger.logsDebug()) {
       operationLogger.debug("onDataUpdate: " + path + " " + message);
     }
-    dataUpdateCount++; // For testing.
 
     List<? extends Event> events;
     try {
@@ -367,7 +322,6 @@ public class Repo implements PersistentConnection.Delegate {
     if (dataLogger.logsDebug()) {
       operationLogger.debug("onRangeMergeUpdate: " + path + " " + merges);
     }
-    dataUpdateCount++; // For testing.
 
     List<RangeMerge> parsedMerges = new ArrayList<>(merges.size());
     for (com.google.firebase.database.connection.RangeMerge merge : merges) {
@@ -610,7 +564,7 @@ public class Repo implements PersistentConnection.Delegate {
     onServerInfoUpdate(Constants.DOT_INFO_AUTHENTICATED, authOk);
   }
 
-  public void onServerInfoUpdate(ChildKey key, Object value) {
+  private void onServerInfoUpdate(ChildKey key, Object value) {
     updateInfo(key, value);
   }
 
@@ -941,8 +895,8 @@ public class Repo implements PersistentConnection.Delegate {
               repo.postEvents(events);
 
               // Finally, run the callbacks
-              for (int i = 0; i < callbacks.size(); ++i) {
-                postEvent(callbacks.get(i));
+              for (Runnable callback : callbacks) {
+                postEvent(callback);
               }
             } else {
               // transactions are no longer sent. Update their status appropriately
@@ -1143,8 +1097,8 @@ public class Repo implements PersistentConnection.Delegate {
     pruneCompletedTransactions(transactionQueueTree);
 
     // Now fire callbacks, now that we're in a good, known state.
-    for (int i = 0; i < callbacks.size(); ++i) {
-      postEvent(callbacks.get(i));
+    for (Runnable callback : callbacks) {
+      postEvent(callback);
     }
 
     // Try to send the transaction result to the server.
@@ -1364,13 +1318,7 @@ public class Repo implements PersistentConnection.Delegate {
 
     @Override
     public int compareTo(TransactionData o) {
-      if (order < o.order) {
-        return -1;
-      } else if (order == o.order) {
-        return 0;
-      } else {
-        return 1;
-      }
+      return Long.compare(order, o.order);
     }
   }
 }
diff --git a/src/main/java/com/google/firebase/database/util/AndroidSupport.java b/src/main/java/com/google/firebase/database/util/AndroidSupport.java
deleted file mode 100644
index 85f5635..0000000
--- a/src/main/java/com/google/firebase/database/util/AndroidSupport.java
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * Copyright 2017 Google Inc.
- *
- * Licensed 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 com.google.firebase.database.util;
-
-public class AndroidSupport {
-
-  private static final boolean IS_ANDROID = checkAndroid();
-
-  private static boolean checkAndroid() {
-    try {
-      Class contextClass = Class.forName("android.app.Activity");
-      return true;
-    } catch (ClassNotFoundException e) {
-      return false;
-    }
-  }
-
-  public static boolean isAndroid() {
-    return IS_ANDROID;
-  }
-}
diff --git a/src/main/java/com/google/firebase/internal/GaeExecutorService.java b/src/main/java/com/google/firebase/internal/GaeExecutorService.java
index 6ee3b16..6d1b323 100644
--- a/src/main/java/com/google/firebase/internal/GaeExecutorService.java
+++ b/src/main/java/com/google/firebase/internal/GaeExecutorService.java
@@ -29,6 +29,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -43,28 +44,34 @@ import java.util.concurrent.atomic.AtomicReference;
  */
 class GaeExecutorService implements ExecutorService {
 
-  private final AtomicReference<ExecutorWrapper> executor = new AtomicReference<>();
+  private final AtomicReference<ExecutorService> executor = new AtomicReference<>();
   private final String threadName;
+  private final ThreadFactory threadFactory;
   private boolean shutdown;
 
   GaeExecutorService(String threadName) {
+    this(threadName, GaeThreadFactory.getInstance());
+  }
+
+  GaeExecutorService(String threadName, ThreadFactory threadFactory) {
     checkArgument(!Strings.isNullOrEmpty(threadName));
     this.threadName = threadName;
+    this.threadFactory = threadFactory;
   }
 
   private ExecutorService ensureExecutorService() {
-    ExecutorWrapper wrapper = executor.get();
-    if (wrapper == null) {
+    ExecutorService executorService = executor.get();
+    if (executorService == null) {
       synchronized (executor) {
         checkState(!shutdown);
-        wrapper = executor.get();
-        if (wrapper == null) {
-          wrapper = new ExecutorWrapper(threadName);
-          executor.compareAndSet(null, wrapper);
+        executorService = executor.get();
+        if (executorService == null) {
+          executorService = newExecutorService(threadFactory, threadName);
+          executor.compareAndSet(null, executorService);
         }
       }
     }
-    return wrapper.getExecutorService();
+    return executorService;
   }
 
   @Override
@@ -110,9 +117,9 @@ class GaeExecutorService implements ExecutorService {
   @Override
   public void shutdown() {
     synchronized (executor) {
-      ExecutorWrapper wrapper = executor.get();
-      if (wrapper != null && !shutdown) {
-        wrapper.getExecutorService().shutdown();
+      ExecutorService executorService = executor.get();
+      if (executorService != null && !shutdown) {
+        executorService.shutdown();
       }
       shutdown = true;
     }
@@ -121,10 +128,10 @@ class GaeExecutorService implements ExecutorService {
   @Override
   public List<Runnable> shutdownNow() {
     synchronized (executor) {
-      ExecutorWrapper wrapper = executor.get();
+      ExecutorService executorService = executor.get();
       List<Runnable> result;
-      if (wrapper != null && !shutdown) {
-        result = wrapper.getExecutorService().shutdownNow();
+      if (executorService != null && !shutdown) {
+        result = executorService.shutdownNow();
       } else {
         result = ImmutableList.of();
       }
@@ -135,17 +142,30 @@ class GaeExecutorService implements ExecutorService {
 
   @Override
   public boolean isShutdown() {
-    return ensureExecutorService().isShutdown();
+    synchronized (executor) {
+      return shutdown;
+    }
   }
 
   @Override
   public boolean isTerminated() {
-    return ensureExecutorService().isTerminated();
+    synchronized (executor) {
+      if (!shutdown) {
+        return false;
+      }
+      ExecutorService executorService = executor.get();
+      return executorService == null || executorService.isTerminated();
+    }
   }
 
   @Override
   public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
-    return ensureExecutorService().awaitTermination(timeout, unit);
+    ExecutorService executorService;
+    synchronized (executor) {
+      executorService = executor.get();
+    }
+    // call await outside the lock
+    return executorService == null || executorService.awaitTermination(timeout, unit);
   }
 
   @Override
@@ -153,31 +173,23 @@ class GaeExecutorService implements ExecutorService {
     ensureExecutorService().execute(command);
   }
 
-  private static class ExecutorWrapper {
-
-    private final ExecutorService executorService;
-
-    ExecutorWrapper(String threadName) {
-      GaeThreadFactory threadFactory = GaeThreadFactory.getInstance();
-      if (threadFactory.isUsingBackgroundThreads()) {
-        // Create a thread pool with long-lived threads if background thread support is available.
-        executorService = new RevivingScheduledExecutor(threadFactory, threadName, true);
-      } else {
-        // Create an executor that creates a new thread for each submitted task, when background
-        // thread support is not available.
-        executorService =
-            new ThreadPoolExecutor(
-                0,
-                Integer.MAX_VALUE,
-                0L,
-                TimeUnit.SECONDS,
-                new SynchronousQueue<Runnable>(),
-                threadFactory);
-      }
-    }
-
-    ExecutorService getExecutorService() {
-      return executorService;
+  private static ExecutorService newExecutorService(
+      ThreadFactory threadFactory, String threadName) {
+    boolean background = threadFactory instanceof GaeThreadFactory
+        && ((GaeThreadFactory) threadFactory).isUsingBackgroundThreads();
+    if (background) {
+      // Create a thread pool with long-lived threads if background thread support is available.
+      return new RevivingScheduledExecutor(threadFactory, threadName, true);
+    } else {
+      // Create an executor that creates a new thread for each submitted task, when background
+      // thread support is not available.
+      return new ThreadPoolExecutor(
+          0,
+          Integer.MAX_VALUE,
+          0L,
+          TimeUnit.SECONDS,
+          new SynchronousQueue<Runnable>(),
+          threadFactory);
     }
   }
 }
diff --git a/src/main/java/com/google/firebase/internal/GaeThreadFactory.java b/src/main/java/com/google/firebase/internal/GaeThreadFactory.java
index 5fd71f0..3791dad 100644
--- a/src/main/java/com/google/firebase/internal/GaeThreadFactory.java
+++ b/src/main/java/com/google/firebase/internal/GaeThreadFactory.java
@@ -87,7 +87,7 @@ public class GaeThreadFactory implements ThreadFactory {
    *
    * @return true if background thread support is available, and false otherwise.
    */
-  public boolean isUsingBackgroundThreads() {
+  boolean isUsingBackgroundThreads() {
     ThreadFactoryWrapper wrapper = threadFactory.get();
     if (wrapper != null) {
       return wrapper.isUsingBackgroundThreads();


Test patch:
diff --git a/src/test/java/com/google/firebase/auth/FirebaseUserManagerTest.java b/src/test/java/com/google/firebase/auth/FirebaseUserManagerTest.java
index 325bf9d..5fe5bc8 100644
--- a/src/test/java/com/google/firebase/auth/FirebaseUserManagerTest.java
+++ b/src/test/java/com/google/firebase/auth/FirebaseUserManagerTest.java
@@ -100,6 +100,22 @@ public class FirebaseUserManagerTest {
     checkRequestHeaders(interceptor);
   }
 
+  @Test
+  public void testGetUserByEmailWithNotFoundError() throws Exception {
+    MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
+    response.setContent(TestUtils.loadResource("getUserError.json"));
+    MockHttpTransport transport = new MockHttpTransport.Builder()
+        .setLowLevelHttpResponse(response)
+        .build();
+    FirebaseUserManager userManager = new FirebaseUserManager(gson, transport, credentials);
+    try {
+      userManager.getUserByEmail("testuser@example.com");
+      fail("No error thrown for invalid response");
+    } catch (FirebaseAuthException e) {
+      assertEquals(FirebaseUserManager.USER_NOT_FOUND_ERROR, e.getErrorCode());
+    }
+  }
+
   @Test
   public void testGetUserByPhoneNumber() throws Exception {
     MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
@@ -115,6 +131,22 @@ public class FirebaseUserManagerTest {
     checkRequestHeaders(interceptor);
   }
 
+  @Test
+  public void testGetUserByPhoneNumberWithNotFoundError() throws Exception {
+    MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
+    response.setContent(TestUtils.loadResource("getUserError.json"));
+    MockHttpTransport transport = new MockHttpTransport.Builder()
+        .setLowLevelHttpResponse(response)
+        .build();
+    FirebaseUserManager userManager = new FirebaseUserManager(gson, transport, credentials);
+    try {
+      userManager.getUserByPhoneNumber("+1234567890");
+      fail("No error thrown for invalid response");
+    } catch (FirebaseAuthException e) {
+      assertEquals(FirebaseUserManager.USER_NOT_FOUND_ERROR, e.getErrorCode());
+    }
+  }
+
   @Test
   public void testListUsers() throws Exception {
     MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
@@ -263,6 +295,44 @@ public class FirebaseUserManagerTest {
 
   @Test
   public void testGetUserHttpError() throws Exception {
+    Map<UserManagerOp, String> operations = ImmutableMap.<UserManagerOp, String>builder()
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.getUserById("testuser");
+          }
+        }, FirebaseUserManager.INTERNAL_ERROR)
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.getUserByEmail("testuser@example.com");
+          }
+        }, FirebaseUserManager.INTERNAL_ERROR)
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.getUserByPhoneNumber("+1234567890");
+          }
+        }, FirebaseUserManager.INTERNAL_ERROR)
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.createUser(new CreateRequest());
+          }
+        }, FirebaseUserManager.USER_CREATE_ERROR)
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.updateUser(new UpdateRequest("test"), Utils.getDefaultJsonFactory());
+          }
+        }, FirebaseUserManager.USER_UPDATE_ERROR)
+        .put(new UserManagerOp() {
+          @Override
+          public void call(FirebaseUserManager userManager) throws Exception {
+            userManager.deleteUser("testuser");
+          }
+        }, FirebaseUserManager.USER_DELETE_ERROR)
+        .build();
     for (int code : ImmutableList.of(302, 400, 401, 404, 500)) {
       MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
       response.setContent("{}");
@@ -271,12 +341,14 @@ public class FirebaseUserManagerTest {
           .setLowLevelHttpResponse(response)
           .build();
       FirebaseUserManager userManager = new FirebaseUserManager(gson, transport, credentials);
-      try {
-        userManager.getUserById("testuser");
-        fail("No error thrown for HTTP error");
-      }  catch (FirebaseAuthException e) {
-        assertTrue(e.getCause() instanceof IOException);
-        assertEquals(FirebaseUserManager.INTERNAL_ERROR, e.getErrorCode());
+      for (Map.Entry<UserManagerOp, String> entry : operations.entrySet()) {
+        try {
+          entry.getKey().call(userManager);
+          fail("No error thrown for HTTP error");
+        }  catch (FirebaseAuthException e) {
+          assertTrue(e.getCause() instanceof IOException);
+          assertEquals(entry.getValue(), e.getErrorCode());
+        }
       }
     }
   }
@@ -696,4 +768,8 @@ public class FirebaseUserManagerTest {
     }
   }
 
+  private interface UserManagerOp {
+    void call(FirebaseUserManager userManager) throws Exception;
+  }
+
 }
diff --git a/src/test/java/com/google/firebase/cloud/StorageClientTest.java b/src/test/java/com/google/firebase/cloud/StorageClientTest.java
index d7b3219..41535f3 100644
--- a/src/test/java/com/google/firebase/cloud/StorageClientTest.java
+++ b/src/test/java/com/google/firebase/cloud/StorageClientTest.java
@@ -17,9 +17,12 @@
 package com.google.firebase.cloud;
 
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.fail;
 
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.storage.Bucket;
+import com.google.cloud.storage.Storage;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
 import com.google.firebase.TestOnlyImplFirebaseTrampolines;
@@ -27,6 +30,7 @@ import com.google.firebase.testing.ServiceAccount;
 import java.io.IOException;
 import org.junit.After;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 public class StorageClientTest {
 
@@ -76,6 +80,7 @@ public class StorageClientTest {
         .setStorageBucket("mock-bucket-name")
         .build());
 
+    assertNotNull(StorageClient.getInstance());
     assertNotNull(StorageClient.getInstance(app));
     app.delete();
     try {
@@ -85,4 +90,40 @@ public class StorageClientTest {
       // ignore
     }
   }
+
+  @Test
+  public void testNonExistingBucket() throws IOException {
+    FirebaseApp app = FirebaseApp.initializeApp(new FirebaseOptions.Builder()
+        .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
+        .setStorageBucket("mock-bucket-name")
+        .build());
+    Storage mockStorage = Mockito.mock(Storage.class);
+    StorageClient client = new StorageClient(app, mockStorage);
+    try {
+      client.bucket();
+      fail("No error thrown for non existing bucket");
+    } catch (IllegalArgumentException expected) {
+      // expected
+    }
+    try {
+      client.bucket("mock-bucket-name");
+      fail("No error thrown for non existing bucket");
+    } catch (IllegalArgumentException expected) {
+      // expected
+    }
+  }
+
+  @Test
+  public void testBucket() throws IOException {
+    FirebaseApp app = FirebaseApp.initializeApp(new FirebaseOptions.Builder()
+        .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
+        .setStorageBucket("mock-bucket-name")
+        .build());
+    Storage mockStorage = Mockito.mock(Storage.class);
+    Bucket mockBucket = Mockito.mock(Bucket.class);
+    Mockito.when(mockStorage.get("mock-bucket-name")).thenReturn(mockBucket);
+    StorageClient client = new StorageClient(app, mockStorage);
+    assertSame(mockBucket, client.bucket());
+    assertSame(mockBucket, client.bucket("mock-bucket-name"));
+  }
 }
diff --git a/src/test/java/com/google/firebase/database/MapperTest.java b/src/test/java/com/google/firebase/database/MapperTest.java
index f6a5892..e7398cf 100644
--- a/src/test/java/com/google/firebase/database/MapperTest.java
+++ b/src/test/java/com/google/firebase/database/MapperTest.java
@@ -34,7 +34,7 @@ import org.junit.Test;
 
 public class MapperTest {
 
-  private static final double EPISLON = 0.00001f;
+  private static final double EPSILON = 0.00001f;
 
   private static <T> T deserialize(String jsonString, Class<T> clazz) {
     Map<String, Object> json = fromSingleQuotedString(jsonString);
@@ -125,14 +125,14 @@ public class MapperTest {
   @Test
   public void primitiveDeserializeDouble() {
     DoubleBean beanDouble = deserialize("{'value': 1.1}", DoubleBean.class);
-    assertEquals(1.1, beanDouble.value, EPISLON);
+    assertEquals(1.1, beanDouble.value, EPSILON);
 
     // Int
     DoubleBean beanInt = deserialize("{'value': 1}", DoubleBean.class);
-    assertEquals(1, beanInt.value, EPISLON);
+    assertEquals(1, beanInt.value, EPSILON);
     // Long
     DoubleBean beanLong = deserialize("{'value': 1234567890123}", DoubleBean.class);
-    assertEquals(1234567890123L, beanLong.value, EPISLON);
+    assertEquals(1234567890123L, beanLong.value, EPSILON);
 
     // Boolean
     try {
@@ -152,14 +152,14 @@ public class MapperTest {
   @Test
   public void primitiveDeserializeFloat() {
     FloatBean beanFloat = deserialize("{'value': 1.1}", FloatBean.class);
-    assertEquals(1.1, beanFloat.value, EPISLON);
+    assertEquals(1.1, beanFloat.value, EPSILON);
 
     // Int
     FloatBean beanInt = deserialize("{'value': 1}", FloatBean.class);
-    assertEquals(1, beanInt.value, EPISLON);
+    assertEquals(1, beanInt.value, EPSILON);
     // Long
     FloatBean beanLong = deserialize("{'value': 1234567890123}", FloatBean.class);
-    assertEquals(Long.valueOf(1234567890123L).floatValue(), beanLong.value, EPISLON);
+    assertEquals(Long.valueOf(1234567890123L).floatValue(), beanLong.value, EPSILON);
 
     // Boolean
     try {
diff --git a/src/test/java/com/google/firebase/database/TestHelpers.java b/src/test/java/com/google/firebase/database/TestHelpers.java
index 548b279..766443d 100644
--- a/src/test/java/com/google/firebase/database/TestHelpers.java
+++ b/src/test/java/com/google/firebase/database/TestHelpers.java
@@ -20,7 +20,10 @@ import static com.cedarsoftware.util.DeepEquals.deepEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.google.common.collect.ImmutableList;
 import com.google.firebase.FirebaseApp;
+import com.google.firebase.database.connection.ConnectionAuthTokenProvider;
+import com.google.firebase.database.connection.ConnectionContext;
 import com.google.firebase.database.core.CoreTestHelpers;
 import com.google.firebase.database.core.DatabaseConfig;
 import com.google.firebase.database.core.Path;
@@ -28,7 +31,10 @@ import com.google.firebase.database.core.Repo;
 import com.google.firebase.database.core.RepoManager;
 import com.google.firebase.database.core.view.QuerySpec;
 import com.google.firebase.database.future.WriteFuture;
+import com.google.firebase.database.logging.DefaultLogger;
+import com.google.firebase.database.logging.Logger.Level;
 import com.google.firebase.database.snapshot.ChildKey;
+import com.google.firebase.database.util.GAuthToken;
 import com.google.firebase.database.util.JsonMapper;
 import com.google.firebase.database.utilities.DefaultRunLoop;
 import com.google.firebase.internal.NonNull;
@@ -43,6 +49,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
@@ -316,6 +323,19 @@ public class TestHelpers {
     return repo;
   }
 
+  public static ConnectionContext newConnectionContext(ScheduledExecutorService executor) {
+    com.google.firebase.database.logging.Logger logger = new DefaultLogger(
+        Level.NONE, ImmutableList.<String>of());
+    ConnectionAuthTokenProvider tokenProvider = new ConnectionAuthTokenProvider() {
+      @Override
+      public void getToken(boolean forceRefresh, GetTokenCallback callback) {
+        callback.onSuccess("gauth|{\"token\":\"test-token\"}");
+      }
+    };
+    return new ConnectionContext(logger, tokenProvider, executor, false, "testVersion",
+        "testUserAgent", Executors.defaultThreadFactory());
+  }
+
   private static class TestExceptionHandler implements UncaughtExceptionHandler {
 
     private final AtomicReference<Throwable> throwable = new AtomicReference<>();
diff --git a/src/test/java/com/google/firebase/database/collection/ArraySortedMapTest.java b/src/test/java/com/google/firebase/database/collection/ArraySortedMapTest.java
index fdaeebe..8f9a0d6 100644
--- a/src/test/java/com/google/firebase/database/collection/ArraySortedMapTest.java
+++ b/src/test/java/com/google/firebase/database/collection/ArraySortedMapTest.java
@@ -83,6 +83,7 @@ public class ArraySortedMapTest {
 
     assertEquals(2, map.size());
     assertEquals(false, map.isEmpty());
+    assertEquals("ArraySortedMap{(a=>1), (b=>null)};", map.toString());
   }
 
   @Test
diff --git a/src/test/java/com/google/firebase/database/collection/ImmutableSortedSetTest.java b/src/test/java/com/google/firebase/database/collection/ImmutableSortedSetTest.java
new file mode 100644
index 0000000..80dce84
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/collection/ImmutableSortedSetTest.java
@@ -0,0 +1,82 @@
+package com.google.firebase.database.collection;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Ordering;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.junit.Test;
+
+public class ImmutableSortedSetTest {
+
+  private static  final List<String> expected = ImmutableList.of("bar", "baz", "foo");
+
+  @Test
+  public void testImmutableSortedSet() {
+    ImmutableSortedSet<String> set = newSet();
+    assertEquals(3, set.size());
+    assertFalse(set.isEmpty());
+    assertEquals("bar", set.getMinEntry());
+    assertEquals("foo", set.getMaxEntry());
+    assertEquals("bar", set.getPredecessorEntry("baz"));
+    assertTrue(set.contains("foo"));
+    assertFalse(set.contains("not found"));
+    assertEquals(set.hashCode(), set.hashCode());
+    assertEquals(set, set);
+
+    ImmutableSortedSet<String> set2 = newSet();
+    assertEquals(set.hashCode(), set2.hashCode());
+    assertEquals(set, set2);
+
+    List<String> original = new ArrayList<>();
+    Collections.addAll(original, "abc", "mno", "pqr");
+    set2 = new ImmutableSortedSet<>(original, Ordering.<String>natural());
+    assertNotEquals(set.hashCode(), set2.hashCode());
+    assertNotEquals(set, set2);
+  }
+
+  @Test
+  public void testForwardIterators() {
+    ImmutableSortedSet<String> set = newSet();
+    Iterator<String> iterator = set.iterator();
+    int i = 0;
+    while (iterator.hasNext()) {
+      assertEquals(expected.get(i++), iterator.next());
+    }
+
+    iterator = set.iteratorFrom("baz");
+    i = 1;
+    while (iterator.hasNext()) {
+      assertEquals(expected.get(i++), iterator.next());
+    }
+  }
+
+  @Test
+  public void testReverseIterators() {
+    ImmutableSortedSet<String> set = newSet();
+    Iterator<String> iterator = set.reverseIterator();
+    int i = 2;
+    while (iterator.hasNext()) {
+      assertEquals(expected.get(i--), iterator.next());
+    }
+
+    iterator = set.reverseIteratorFrom("baz");
+    i = 1;
+    while (iterator.hasNext()) {
+      assertEquals(expected.get(i--), iterator.next());
+    }
+  }
+
+  private ImmutableSortedSet<String> newSet() {
+    List<String> original = new ArrayList<>();
+    Collections.addAll(original, "foo", "bar", "baz");
+    return new ImmutableSortedSet<>(original, Ordering.<String>natural());
+  }
+
+}
diff --git a/src/test/java/com/google/firebase/database/connection/ConnectionTest.java b/src/test/java/com/google/firebase/database/connection/ConnectionTest.java
new file mode 100644
index 0000000..56ff9a9
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/connection/ConnectionTest.java
@@ -0,0 +1,202 @@
+/*
+ * Copyright 2017 Google Inc.
+ *
+ * Licensed 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 com.google.firebase.database.connection;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.firebase.database.TestHelpers;
+import com.google.firebase.database.connection.Connection.Delegate;
+import com.google.firebase.database.connection.Connection.DisconnectReason;
+import com.google.firebase.database.connection.Connection.WebsocketConnectionFactory;
+import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import org.junit.AfterClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class ConnectionTest {
+
+  private static final ScheduledExecutorService executor =
+      Executors.newSingleThreadScheduledExecutor();
+
+  @AfterClass
+  public static void tearDownClass() {
+    executor.shutdownNow();
+  }
+
+  @Test
+  public void testOpen() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).open();
+  }
+
+  @Test
+  public void testClose() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+    connFactory.connection.close();
+
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+  }
+
+  @Test
+  public void testCloseWithReason() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+    connFactory.connection.close(DisconnectReason.SERVER_RESET);
+
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(
+        DisconnectReason.SERVER_RESET);
+  }
+
+  @Test
+  public void testSendRequest() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    // should not be sent, since not connected
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    connFactory.connection.sendRequest(data, true);
+    Mockito.verify(connFactory.wsConn, Mockito.never()).send(Mockito.<String, Object>anyMap());
+
+    Map<String, Object> handshake = ImmutableMap.<String, Object>of(
+        "t", "c",
+        "d", ImmutableMap.of("t", "h", "d", ImmutableMap.of(
+            "ts", System.currentTimeMillis())));
+    connFactory.connection.onMessage(handshake);
+
+    Map<String, Object> expected = ImmutableMap.<String, Object>of("t", "d", "d", data);
+    connFactory.connection.sendRequest(data, true);
+    connFactory.connection.sendRequest(data, false);
+    Mockito.verify(connFactory.wsConn, Mockito.times(2)).send(expected);
+  }
+
+  @Test
+  public void testOnDataMessage() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    Map<String, Object> incoming = ImmutableMap.<String, Object>of(
+        "t", "d",
+        "d", data);
+    connFactory.connection.onMessage(incoming);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDataMessage(data);
+  }
+
+  @Test
+  public void testOnDataMessageParseError() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    Mockito.doThrow(ClassCastException.class).when(connFactory.delegate).onDataMessage(data);
+
+    Map<String, Object> incoming = ImmutableMap.<String, Object>of(
+        "t", "d",
+        "d", data);
+    connFactory.connection.onMessage(incoming);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testOnInvalidMessage() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    Map<String, Object> incoming = ImmutableMap.<String, Object>of("d", data);
+    connFactory.connection.onMessage(incoming);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testOnShutdown() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    Map<String, Object> command = ImmutableMap.<String, Object>of(
+        "t", "c",
+        "d", ImmutableMap.of("t", "s", "d", "shutdown_reason"));
+    connFactory.connection.onMessage(command);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onKill("shutdown_reason");
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testOnReset() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    Map<String, Object> command = ImmutableMap.<String, Object>of(
+        "t", "c",
+        "d", ImmutableMap.of("t", "r", "d", "reset_host"));
+    connFactory.connection.onMessage(command);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onCacheHost("reset_host");
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(
+        DisconnectReason.SERVER_RESET);
+    Mockito.verify(connFactory.wsConn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testOnDisconnect() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    connFactory.connection.onDisconnect(true);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+    Mockito.verify(connFactory.wsConn, Mockito.never()).close();
+  }
+
+  @Test
+  public void testOnDisconnectNeverConnected() {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.connection.open();
+
+    connFactory.connection.onDisconnect(false);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect(DisconnectReason.OTHER);
+    Mockito.verify(connFactory.wsConn, Mockito.never()).close();
+  }
+
+  private static class MockConnectionFactory implements WebsocketConnectionFactory {
+
+    private final Delegate delegate = Mockito.mock(Delegate.class);
+    private final WebsocketConnection wsConn = Mockito.mock(WebsocketConnection.class);
+
+    private final Connection connection;
+
+    MockConnectionFactory() {
+      this.connection = new Connection(
+          TestHelpers.newConnectionContext(executor),
+          Mockito.mock(HostInfo.class),
+          delegate,
+          this);
+    }
+
+    @Override
+    public WebsocketConnection newConnection(WebsocketConnection.Delegate delegate) {
+      return wsConn;
+    }
+  }
+
+}
diff --git a/src/test/java/com/google/firebase/database/connection/ConnectionUtilsTest.java b/src/test/java/com/google/firebase/database/connection/ConnectionUtilsTest.java
new file mode 100644
index 0000000..5f892bf
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/connection/ConnectionUtilsTest.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2017 Google Inc.
+ *
+ * Licensed 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 com.google.firebase.database.connection;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+
+import com.google.common.collect.ImmutableList;
+import org.junit.Test;
+
+public class ConnectionUtilsTest {
+
+  @Test
+  public void testStringToPath() {
+    assertEquals(ImmutableList.of(), ConnectionUtils.stringToPath(""));
+    assertEquals(ImmutableList.of(), ConnectionUtils.stringToPath("/"));
+    assertEquals(ImmutableList.of("foo"), ConnectionUtils.stringToPath("/foo"));
+    assertEquals(ImmutableList.of("foo", "bar"), ConnectionUtils.stringToPath("/foo/bar"));
+    assertEquals(ImmutableList.of("foo", "bar"), ConnectionUtils.stringToPath("/foo//bar"));
+  }
+
+  @Test
+  public void testPathToString() {
+    assertEquals("/", ConnectionUtils.pathToString(ImmutableList.<String>of()));
+    assertEquals("foo", ConnectionUtils.pathToString(ImmutableList.of("foo")));
+    assertEquals("foo/bar", ConnectionUtils.pathToString(ImmutableList.of("foo", "bar")));
+  }
+
+  @Test
+  public void testLongFromObject() {
+    assertEquals(new Long(10), ConnectionUtils.longFromObject(10L));
+    assertEquals(new Long(10), ConnectionUtils.longFromObject(new Long(10)));
+    assertEquals(new Long(10), ConnectionUtils.longFromObject(10));
+    assertNull(ConnectionUtils.longFromObject("foo"));
+  }
+
+  @Test
+  public void testHardAssert() {
+    ConnectionUtils.hardAssert(true);
+    try {
+      ConnectionUtils.hardAssert(false);
+      fail("No error thrown for failed hardAssert");
+    } catch (AssertionError e) {
+      assertEquals("hardAssert failed: ", e.getMessage());
+    }
+  }
+
+  @Test
+  public void testHardAssertWithMessage() {
+    ConnectionUtils.hardAssert(true, "foo %s", "bar");
+    try {
+      ConnectionUtils.hardAssert(false,"foo %s", "bar");
+      fail("No error thrown for failed hardAssert");
+    } catch (AssertionError e) {
+      assertEquals("hardAssert failed: foo bar", e.getMessage());
+    }
+  }
+
+}
diff --git a/src/test/java/com/google/firebase/database/connection/PersistentConnectionTest.java b/src/test/java/com/google/firebase/database/connection/PersistentConnectionTest.java
new file mode 100644
index 0000000..209ceda
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/connection/PersistentConnectionTest.java
@@ -0,0 +1,305 @@
+package com.google.firebase.database.connection;
+
+import static com.google.firebase.database.TestHelpers.waitFor;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.firebase.database.TestHelpers;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.Semaphore;
+
+import org.junit.AfterClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+public class PersistentConnectionTest {
+
+  private static final ScheduledExecutorService executor =
+      Executors.newSingleThreadScheduledExecutor();
+
+  @AfterClass
+  public static void tearDownClass() {
+    executor.shutdownNow();
+  }
+
+  @Test
+  public void testOnReady() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onServerInfoUpdate(
+        Mockito.<String, Object>anyMap());
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onConnect();
+    assertEquals(2, connFactory.outgoing.size());
+    assertEquals("s", connFactory.outgoing.get(0).getAction());
+    assertEquals("gauth", connFactory.outgoing.get(1).getAction());
+    assertTrue(connFactory.outgoing.get(1).sensitive);
+  }
+
+  @Test
+  public void testAuthSuccess() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 1, "b", ImmutableMap.of("s", "ok")));
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onAuthStatus(true);
+  }
+
+  @Test
+  public void testAuthFailure() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "r", 1, "b", ImmutableMap.of("s", "not_ok")));
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onAuthStatus(false);
+  }
+
+  @Test
+  public void testUpgradeAuth() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.refreshAuthToken("new-token");
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "r", 1, "b", ImmutableMap.of("s", "not_ok")));
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("auth", connFactory.outgoing.get(2).getAction());
+  }
+
+  @Test
+  public void testUnauth() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.refreshAuthToken(null);
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "r", 1, "b", ImmutableMap.of("s", "not_ok")));
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("unauth", connFactory.outgoing.get(2).getAction());
+  }
+
+  @Test
+  public void testOnDataMessage() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "a", "d", "b", ImmutableMap.of("p", "foo", "d", "test")));
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDataUpdate(
+        ImmutableList.of("foo"), "test", false, null);
+
+    List<Map<String, Object>> ranges = ImmutableList.<Map<String, Object>>of(
+        ImmutableMap.<String, Object>of("s", "start", "e", "end", "m", "data"));
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "a", "rm", "b", ImmutableMap.of("p", "foo", "d", ranges)));
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onRangeMergeUpdate(
+        Mockito.<String>anyList(), Mockito.<RangeMerge>anyList(), Mockito.nullable(Long.class));
+
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "a", "ac", "b", ImmutableMap.of("s", "status", "d", "reason")));
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onAuthStatus(false);
+  }
+
+  @Test
+  public void testListen() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    ListenHashProvider hash = Mockito.mock(ListenHashProvider.class);
+    Mockito.when(hash.getSimpleHash()).thenReturn("simpleHash");
+
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.listen(
+        ImmutableList.of("listen"), ImmutableMap.<String, Object>of(), hash, null, callback);
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("q", connFactory.outgoing.get(2).getAction());
+
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "r", 2, "b", ImmutableMap.of("s", "ok", "d", ImmutableMap.of())));
+    Mockito.verify(callback, Mockito.times(1)).onRequestResult(null, null);
+
+    connFactory.persistentConn.onDataMessage(ImmutableMap.<String, Object>of(
+        "a", "c", "b", ImmutableMap.of("p", "listen")));
+    Mockito.verify(callback, Mockito.times(1)).onRequestResult("permission_denied", null);
+  }
+
+  @Test
+  public void testUnlisten() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    ListenHashProvider hash = Mockito.mock(ListenHashProvider.class);
+    Mockito.when(hash.getSimpleHash()).thenReturn("simpleHash");
+
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.listen(
+        ImmutableList.of("listen"), ImmutableMap.<String, Object>of(), hash, null, callback);
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("q", connFactory.outgoing.get(2).getAction());
+
+    connFactory.persistentConn.unlisten(
+        ImmutableList.of("listen"), ImmutableMap.<String, Object>of());
+    assertEquals(4, connFactory.outgoing.size());
+    assertEquals("n", connFactory.outgoing.get(3).getAction());
+  }
+
+  @Test
+  public void testPut() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.put(ImmutableList.of("put"), "testData", callback);
+    assertEquals(2, connFactory.outgoing.size());
+
+    // Write should not be sent until auth is successful
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 1, "b", ImmutableMap.of("s", "ok")));
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("p", connFactory.outgoing.get(2).getAction());
+
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 2, "b", ImmutableMap.of("s", "ok")));
+    Mockito.verify(callback, Mockito.times(1)).onRequestResult(null, null);
+  }
+
+  @Test
+  public void testOnDisconnectPut() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.onDisconnectPut(ImmutableList.of("put"), "testData", callback);
+    assertEquals(2, connFactory.outgoing.size());
+
+    // Write should not be sent until auth is successful
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 1, "b", ImmutableMap.of("s", "ok")));
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("o", connFactory.outgoing.get(2).getAction());
+
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 2, "b", ImmutableMap.of("s", "ok")));
+    Mockito.verify(callback, Mockito.times(1)).onRequestResult(null, null);
+  }
+
+  @Test
+  public void testOnDisconnectMerge() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 1, "b", ImmutableMap.of("s", "ok")));
+
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.onDisconnectMerge(ImmutableList.of("put"),
+        ImmutableMap.<String, Object>of("key", "value"), callback);
+
+    // Write should not be sent until auth us successful
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("om", connFactory.outgoing.get(2).getAction());
+  }
+
+  @Test
+  public void testOnDisconnectCancel() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+    connFactory.persistentConn.onDataMessage(
+        ImmutableMap.<String, Object>of("r", 1, "b", ImmutableMap.of("s", "ok")));
+
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.onDisconnectCancel(ImmutableList.of("put"), callback);
+
+    // Write should not be sent until auth us successful
+    assertEquals(3, connFactory.outgoing.size());
+    assertEquals("oc", connFactory.outgoing.get(2).getAction());
+  }
+
+  @Test
+  public void testOnKill() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onKill("some_reason");
+    Mockito.verify(connFactory.conn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testShutdown() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.shutdown();
+    Mockito.verify(connFactory.conn, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testOnDisconnect() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onDisconnect(Connection.DisconnectReason.OTHER);
+    Mockito.verify(connFactory.delegate, Mockito.times(1)).onDisconnect();
+  }
+
+  @Test
+  public void testPurgeWrites() throws InterruptedException {
+    MockConnectionFactory connFactory = new MockConnectionFactory();
+    connFactory.persistentConn.onReady(System.currentTimeMillis(), "last-session-id");
+
+    RequestResultCallback callback = Mockito.mock(RequestResultCallback.class);
+    connFactory.persistentConn.put(ImmutableList.of("put"), "testData", callback);
+    connFactory.persistentConn.onDisconnectPut(ImmutableList.of("put"), "testData", callback);
+    assertEquals(2, connFactory.outgoing.size());
+
+    // Write should not be sent until auth is successful
+    connFactory.persistentConn.purgeOutstandingWrites();
+    Mockito.verify(callback, Mockito.times(2)).onRequestResult("write_canceled", null);
+  }
+
+  private static class MockConnectionFactory implements PersistentConnectionImpl.ConnectionFactory {
+
+    private final Connection conn = Mockito.mock(Connection.class);
+    private final PersistentConnection.Delegate delegate = Mockito.mock(
+        PersistentConnection.Delegate.class);
+    private final List<OutgoingMessage> outgoing = new ArrayList<>();
+    private final Semaphore connected = new Semaphore(0);
+
+    private final PersistentConnectionImpl persistentConn;
+
+    MockConnectionFactory() throws InterruptedException {
+      Mockito.doAnswer(new Answer() {
+        @Override
+        public Object answer(InvocationOnMock invocation) throws Throwable {
+          outgoing.add(new OutgoingMessage(
+              (Map) invocation.getArgument(0), (boolean) invocation.getArgument(1)));
+          return null;
+        }
+      }).when(conn).sendRequest(Mockito.<String, Object>anyMap(), Mockito.anyBoolean());
+
+      Mockito.doAnswer(new Answer() {
+        @Override
+        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
+          connected.release();
+          return null;
+        }
+      }).when(conn).open();
+
+      persistentConn = new PersistentConnectionImpl(
+          TestHelpers.newConnectionContext(executor), null, delegate, this);
+      persistentConn.initialize();
+      waitFor(connected);
+    }
+
+    @Override
+    public Connection newConnection(PersistentConnectionImpl delegate) {
+      return conn;
+    }
+  }
+
+  private static class OutgoingMessage {
+    private final Map payload;
+    private final boolean sensitive;
+
+    OutgoingMessage(Map payload, boolean sensitive) {
+      this.payload = payload;
+      this.sensitive = sensitive;
+    }
+
+    String getAction() {
+      return (String) this.payload.get("a");
+    }
+  }
+}
diff --git a/src/test/java/com/google/firebase/database/connection/WebsocketConnectionTest.java b/src/test/java/com/google/firebase/database/connection/WebsocketConnectionTest.java
new file mode 100644
index 0000000..896d4fe
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/connection/WebsocketConnectionTest.java
@@ -0,0 +1,167 @@
+/*
+ * Copyright 2017 Google Inc.
+ *
+ * Licensed 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 com.google.firebase.database.connection;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.firebase.database.TestHelpers;
+import com.google.firebase.database.connection.WebsocketConnection.Delegate;
+import com.google.firebase.database.connection.WebsocketConnection.WSClient;
+import com.google.firebase.database.connection.WebsocketConnection.WSClientEventHandler;
+import com.google.firebase.database.connection.WebsocketConnection.WSClientFactory;
+import java.io.IOException;
+import java.util.concurrent.ScheduledExecutorService;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+public class WebsocketConnectionTest {
+
+  private static final ScheduledExecutorService executor =
+      Mockito.spy(ScheduledExecutorService.class);
+
+  @BeforeClass
+  public static void setUpClass() {
+    // Rig the execute() method to execute the Runnable on the calling thread.
+    Mockito.doAnswer(new Answer() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        Runnable runnable = invocation.getArgument(0);
+        runnable.run();
+        return null;
+      }
+    }).when(executor).execute(Mockito.any(Runnable.class));
+  }
+
+  @AfterClass
+  public static void tearDownClass() {
+    executor.shutdownNow();
+  }
+
+  @Test
+  public void testOpen() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.conn.open();
+    Mockito.verify(clientFactory.client, Mockito.times(1)).connect();
+  }
+
+  @Test
+  public void testClose() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.conn.close();
+    Mockito.verify(clientFactory.client, Mockito.times(1)).close();
+  }
+
+  @Test
+  public void testSendSingleFrame() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.conn.send(ImmutableMap.<String, Object>of("key", "value"));
+    Mockito.verify(clientFactory.client, Mockito.times(1)).send("{\"key\":\"value\"}");
+  }
+
+  @Test
+  public void testSendMultipleFrames() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    StringBuilder longString = new StringBuilder();
+    for (int i = 0; i < 20000; i++) {
+      longString.append("a");
+    }
+    clientFactory.conn.send(ImmutableMap.<String, Object>of("key", longString.toString()));
+
+    String serialized = String.format("{\"key\":\"%s\"}", longString.toString());
+    Mockito.verify(clientFactory.client, Mockito.times(1)).send("2");
+    Mockito.verify(clientFactory.client, Mockito.times(1)).send(serialized.substring(0, 16384));
+    Mockito.verify(clientFactory.client, Mockito.times(1)).send(serialized.substring(16384));
+  }
+
+  @Test
+  public void testReceiveSingleFrame() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.eventHandler.onMessage("{\"key\":\"value\"}");
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onMessage(
+        ImmutableMap.<String, Object>of("key", "value"));
+  }
+
+  @Test
+  public void testReceiveMultipleFrames() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.eventHandler.onMessage("2");
+    clientFactory.eventHandler.onMessage("{\"key1\":\"value1\",");
+    clientFactory.eventHandler.onMessage("\"key2\":\"value2\"}");
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onMessage(
+        ImmutableMap.<String, Object>of("key1", "value1", "key2", "value2"));
+  }
+
+  @Test
+  public void testIncomingMessageIOError() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    Mockito.doThrow(IOException.class).when(clientFactory.delegate).onMessage(data);
+    clientFactory.eventHandler.onMessage("{\"key\":\"value\"}");
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onMessage(data);
+    Mockito.verify(clientFactory.client, Mockito.times(1)).close();
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onDisconnect(false);
+  }
+
+  @Test
+  public void testIncomingMessageClassCastError() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    ImmutableMap<String, Object> data = ImmutableMap.<String, Object>of("key", "value");
+    Mockito.doThrow(ClassCastException.class).when(clientFactory.delegate).onMessage(data);
+    clientFactory.eventHandler.onMessage("{\"key\":\"value\"}");
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onMessage(data);
+    Mockito.verify(clientFactory.client, Mockito.times(1)).close();
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onDisconnect(false);
+  }
+
+  @Test
+  public void testOnClose() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.eventHandler.onClose();
+    Mockito.verify(clientFactory.client, Mockito.times(1)).close();
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onDisconnect(false);
+  }
+
+  @Test
+  public void testOnError() {
+    MockClientFactory clientFactory = new MockClientFactory();
+    clientFactory.eventHandler.onError(new Exception("test"));
+    Mockito.verify(clientFactory.client, Mockito.times(1)).close();
+    Mockito.verify(clientFactory.delegate, Mockito.times(1)).onDisconnect(false);
+  }
+
+  private static class MockClientFactory implements WSClientFactory {
+
+    private final WSClient client = Mockito.mock(WSClient.class);
+    private final Delegate delegate = Mockito.mock(Delegate.class);
+    private final WebsocketConnection conn;
+    private WSClientEventHandler eventHandler;
+
+    MockClientFactory() {
+      ConnectionContext context = TestHelpers.newConnectionContext(executor);
+      this.conn = new WebsocketConnection(context, delegate, this);
+    }
+
+    @Override
+    public WSClient newClient(WSClientEventHandler delegate) {
+      this.eventHandler = delegate;
+      return client;
+    }
+  }
+}
diff --git a/src/test/java/com/google/firebase/database/core/GaePlatformTest.java b/src/test/java/com/google/firebase/database/core/GaePlatformTest.java
new file mode 100644
index 0000000..032de1b
--- /dev/null
+++ b/src/test/java/com/google/firebase/database/core/GaePlatformTest.java
@@ -0,0 +1,95 @@
+package com.google.firebase.database.core;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.firebase.FirebaseApp;
+import com.google.firebase.FirebaseOptions;
+import com.google.firebase.TestOnlyImplFirebaseTrampolines;
+import com.google.firebase.ThreadManager;
+import com.google.firebase.database.FirebaseDatabase;
+import com.google.firebase.database.logging.Logger;
+import com.google.firebase.internal.GaeThreadFactory;
+import com.google.firebase.internal.NonNull;
+import com.google.firebase.testing.ServiceAccount;
+import com.google.firebase.testing.TestUtils;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class GaePlatformTest {
+
+  @Test
+  public void testIsActive() {
+    assertEquals(GaeThreadFactory.isAvailable(), GaePlatform.isActive());
+  }
+
+  @Test
+  public void testGaePlatform() {
+    final AtomicInteger count = new AtomicInteger(0);
+    final ThreadManager threadManager = new ThreadManager() {
+      @Override
+      protected ExecutorService getExecutor(@NonNull FirebaseApp app) {
+        return Executors.newSingleThreadExecutor();
+      }
+
+      @Override
+      protected void releaseExecutor(@NonNull FirebaseApp app,
+                                     @NonNull ExecutorService executor) {
+      }
+
+      @Override
+      protected ThreadFactory getThreadFactory() {
+        count.incrementAndGet();
+        return Executors.defaultThreadFactory();
+      }
+    };
+
+    FirebaseOptions options = new FirebaseOptions.Builder()
+        .setCredentials(TestUtils.getCertCredential(ServiceAccount.EDITOR.asStream()))
+        .setThreadManager(threadManager)
+        .build();
+    FirebaseApp app = FirebaseApp.initializeApp(options, "gaeApp");
+    try {
+      GaePlatform platform = new GaePlatform(app);
+      Context ctx = Mockito.mock(Context.class);
+      assertNotNull(platform.newLogger(ctx, Logger.Level.DEBUG, ImmutableList.<String>of()));
+
+      assertNotNull(platform.newEventTarget(ctx));
+      assertEquals(1, count.get());
+
+      assertNotNull(platform.newRunLoop(ctx));
+      assertEquals(2, count.get());
+
+      AuthTokenProvider authTokenProvider = platform.newAuthTokenProvider(
+          Mockito.mock(ScheduledExecutorService.class));
+      assertTrue(authTokenProvider instanceof JvmAuthTokenProvider);
+
+      assertEquals("AppEngine/AdminJava", platform.getUserAgent(ctx));
+      assertEquals("gae-" + FirebaseDatabase.getSdkVersion(), platform.getPlatformVersion());
+      assertNull(platform.createPersistenceManager(ctx, "test"));
+
+      ThreadInitializer threadInitializer = platform.getThreadInitializer();
+      Thread t = new Thread();
+      threadInitializer.setName(t, "test-name");
+      threadInitializer.setDaemon(t, true);
+      threadInitializer.setUncaughtExceptionHandler(t, Mockito.mock(
+          Thread.UncaughtExceptionHandler.class));
+      assertNotEquals("test-name", t.getName());
+      assertFalse(t.isDaemon());
+      assertNotNull(t.getUncaughtExceptionHandler());
+
+    } finally {
+      TestOnlyImplFirebaseTrampolines.clearInstancesForTest();
+    }
+  }
+}
diff --git a/src/test/java/com/google/firebase/database/core/RepoTest.java b/src/test/java/com/google/firebase/database/core/RepoTest.java
index 89f93c5..ec9696e 100644
--- a/src/test/java/com/google/firebase/database/core/RepoTest.java
+++ b/src/test/java/com/google/firebase/database/core/RepoTest.java
@@ -1,6 +1,7 @@
 package com.google.firebase.database.core;
 
 import static com.google.firebase.database.TestHelpers.waitFor;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -34,13 +35,12 @@ import com.google.firebase.database.snapshot.NodeUtilities;
 import com.google.firebase.database.snapshot.ValueIndex;
 import com.google.firebase.testing.ServiceAccount;
 
-import com.google.firebase.testing.TestUtils;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.AfterClass;
@@ -56,6 +56,7 @@ public class RepoTest {
   private static final List<String> FAILURE = ImmutableList.of("failure");
 
   private static DatabaseConfig config;
+  private static PersistentConnection connection;
 
   @BeforeClass
   public static void setUpClass() throws IOException {
@@ -81,10 +82,23 @@ public class RepoTest {
       }
     });
 
-    PersistentConnection connection = mockConnection();
+    connection = mockConnection();
     Mockito.when(config.newPersistentConnection(
         Mockito.any(HostInfo.class), Mockito.any(PersistentConnection.Delegate.class)))
         .thenReturn(connection);
+    Mockito.doAnswer(new Answer() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        RequestResultCallback callback = invocation.getArgument(3);
+        callback.onRequestResult(null, null);
+        return null;
+      }
+    })
+        .when(connection).compareAndPut(
+        Mockito.eq(ImmutableList.of("txn")),
+        Mockito.any(),
+        Mockito.anyString(),
+        Mockito.any(RequestResultCallback.class));
   }
 
   @AfterClass
@@ -346,6 +360,110 @@ public class RepoTest {
     assertEquals(update, events.get(0).getValue());
   }
 
+  @Test
+  public void testTransactionError() throws InterruptedException {
+    Mockito.doAnswer(new Answer() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        RequestResultCallback callback = invocation.getArgument(3);
+        callback.onRequestResult("test error", "test error message");
+        return null;
+      }
+    })
+        .when(connection).compareAndPut(
+        Mockito.eq(ImmutableList.of("txn_error")),
+        Mockito.any(),
+        Mockito.anyString(),
+        Mockito.any(RequestResultCallback.class));
+
+    final Repo repo = newRepo();
+    final List<DataSnapshot> events = new ArrayList<>();
+    final Path path = new Path("/txn_error");
+    QuerySpec spec = new QuerySpec(path, QueryParams.DEFAULT_PARAMS);
+    addCallback(repo, new ValueEventRegistration(repo, newValueEventListener(events), spec));
+
+    final Semaphore semaphore = new Semaphore(0);
+    final Map<String, Object> update = ImmutableMap.<String, Object>of("key", "value");
+    final AtomicInteger count = new AtomicInteger(0);
+    repo.scheduleNow(new Runnable() {
+      @Override
+      public void run() {
+        repo.startTransaction(path, new Transaction.Handler() {
+          @Override
+          public Transaction.Result doTransaction(MutableData currentData) {
+            count.incrementAndGet();
+            currentData.setValue(update);
+            return Transaction.success(currentData);
+          }
+
+          @Override
+          public void onComplete(DatabaseError error, boolean committed, DataSnapshot currentData) {
+            assertEquals(1, count.get());
+            assertFalse(committed);
+            semaphore.release();
+          }
+        }, true);
+        semaphore.release();
+      }
+    });
+    waitFor(semaphore, 2);
+    assertEquals(1, events.size());
+    assertNotNull(events.get(0));
+    assertEquals(update, events.get(0).getValue());
+  }
+
+  @Test
+  public void testTransactionDataStaleError() throws InterruptedException {
+    Mockito.doAnswer(new Answer() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        RequestResultCallback callback = invocation.getArgument(3);
+        callback.onRequestResult("datastale", "test error message");
+        return null;
+      }
+    })
+        .when(connection).compareAndPut(
+        Mockito.eq(ImmutableList.of("txn_stale_error")),
+        Mockito.any(),
+        Mockito.anyString(),
+        Mockito.any(RequestResultCallback.class));
+
+    final Repo repo = newRepo();
+    final List<DataSnapshot> events = new ArrayList<>();
+    final Path path = new Path("/txn_stale_error");
+    QuerySpec spec = new QuerySpec(path, QueryParams.DEFAULT_PARAMS);
+    addCallback(repo, new ValueEventRegistration(repo, newValueEventListener(events), spec));
+
+    final Semaphore semaphore = new Semaphore(0);
+    final Map<String, Object> update = ImmutableMap.<String, Object>of("key", "value");
+    final AtomicInteger count = new AtomicInteger(0);
+    repo.scheduleNow(new Runnable() {
+      @Override
+      public void run() {
+        repo.startTransaction(path, new Transaction.Handler() {
+          @Override
+          public Transaction.Result doTransaction(MutableData currentData) {
+            count.incrementAndGet();
+            currentData.setValue(update);
+            return Transaction.success(currentData);
+          }
+
+          @Override
+          public void onComplete(DatabaseError error, boolean committed, DataSnapshot currentData) {
+            assertEquals(25, count.get());
+            assertFalse(committed);
+            semaphore.release();
+          }
+        }, true);
+        semaphore.release();
+      }
+    });
+    waitFor(semaphore, 2);
+    assertEquals(1, events.size());
+    assertNotNull(events.get(0));
+    assertEquals(update, events.get(0).getValue());
+  }
+
   @Test
   public void testTransactionAbort() throws InterruptedException {
     final Repo repo = newRepo();
diff --git a/src/test/java/com/google/firebase/database/snapshot/NodeTest.java b/src/test/java/com/google/firebase/database/snapshot/NodeTest.java
index 22e6d6d..f94cea9 100644
--- a/src/test/java/com/google/firebase/database/snapshot/NodeTest.java
+++ b/src/test/java/com/google/firebase/database/snapshot/NodeTest.java
@@ -18,10 +18,19 @@ package com.google.firebase.database.snapshot;
 
 import static com.google.firebase.database.snapshot.NodeUtilities.NodeFromJSON;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.firebase.database.MapBuilder;
+import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.core.Path;
+import com.google.firebase.database.snapshot.LeafNode.LeafType;
+import com.google.firebase.database.snapshot.Node.HashVersion;
 import java.util.Map;
 import org.junit.Test;
 
@@ -180,4 +189,65 @@ public class NodeTest {
         NodeFromJSON(new MapBuilder().put("dummy-node", null).put(".priority", "prio").build());
     assertTrue(empty2.getPriority().isEmpty());
   }
+
+  @Test
+  public void testEmptyChildrenNode() {
+    ChildrenNode node = new ChildrenNode();
+    assertTrue(node.isEmpty());
+    assertNull(node.getValue());
+    assertEquals("{ }", node.toString());
+    assertEquals(0, node.hashCode());
+    assertEquals(0, node.compareTo(new ChildrenNode()));
+  }
+
+  @Test
+  public void testEmptyChildrenNodeWithPriority() {
+    ImmutableSortedMap<ChildKey, Node> map = ImmutableSortedMap.Builder.fromMap(
+        ImmutableMap.<ChildKey, Node>of(), null);
+    try {
+      new ChildrenNode(map, PriorityUtilities.parsePriority(1));
+      fail("No error thrown for empty children node");
+    } catch (IllegalArgumentException expected) {
+      // expected
+    }
+  }
+
+  @Test
+  public void testChildrenNode() {
+    ImmutableMap<ChildKey, Node> source = ImmutableMap.of(
+        ChildKey.fromString("foo"), NodeUtilities.NodeFromJSON("value1"),
+        ChildKey.fromString("bar"), NodeUtilities.NodeFromJSON("value2")
+    );
+    ImmutableSortedMap<ChildKey, Node> map = ImmutableSortedMap.Builder.fromMap(source, null);
+    ChildrenNode node = new ChildrenNode(map, PriorityUtilities.parsePriority(1));
+    assertFalse(node.isEmpty());
+    assertEquals(ImmutableMap.of("foo", "value1", "bar", "value2"), node.getValue());
+    assertEquals(
+        "{\n"
+        + "  bar=value2\n"
+        + "  foo=value1\n"
+        + "  .priority=1.0\n"
+        + "}", node.toString());
+
+    ChildrenNode other = new ChildrenNode(map, PriorityUtilities.parsePriority(1));
+    assertEquals(node, other);
+    assertEquals(node.hashCode(), other.hashCode());
+  }
+
+  @Test
+  public void testDeferredValueNode() {
+    Map<Object, Object> map = ImmutableMap.<Object, Object>of("foo", "bar", "int", 5);
+    DeferredValueNode node = new DeferredValueNode(map, PriorityUtilities.parsePriority(1));
+    assertEquals(map, node.getValue());
+    assertEquals(LeafType.DeferredValue, node.getLeafType());
+    assertNotNull(node.getHashRepresentation(HashVersion.V1));
+
+    DeferredValueNode otherNode = new DeferredValueNode(
+        ImmutableMap.<Object, Object>of("other", "value"),
+        PriorityUtilities.parsePriority(2));
+    assertEquals(0, node.compareLeafValues(otherNode));
+
+    assertNotEquals(node, otherNode);
+    assertNotEquals(node.hashCode(), otherNode.hashCode());
+  }
 }
diff --git a/src/test/java/com/google/firebase/internal/GaeExecutorServiceTest.java b/src/test/java/com/google/firebase/internal/GaeExecutorServiceTest.java
new file mode 100644
index 0000000..48b2678
--- /dev/null
+++ b/src/test/java/com/google/firebase/internal/GaeExecutorServiceTest.java
@@ -0,0 +1,129 @@
+/*
+ * Copyright 2017 Google Inc.
+ *
+ * Licensed 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 com.google.firebase.internal;
+
+import static com.google.firebase.database.TestHelpers.waitFor;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.firebase.testing.TestUtils;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.Test;
+
+public class GaeExecutorServiceTest {
+
+  @Test
+  public void testShutdownBeforeUse() throws InterruptedException {
+    CountingThreadFactory threadFactory = new CountingThreadFactory();
+    GaeExecutorService executorService = new GaeExecutorService("test", threadFactory);
+    assertFalse(executorService.isShutdown());
+    assertFalse(executorService.isTerminated());
+
+    assertEquals(ImmutableList.of(), executorService.shutdownNow());
+    assertTrue(executorService.isShutdown());
+    assertTrue(executorService.isTerminated());
+    assertTrue(executorService.awaitTermination(1, TimeUnit.SECONDS));
+    assertEquals(0, threadFactory.counter.get());
+
+    executorService = new GaeExecutorService("test", threadFactory);
+    assertFalse(executorService.isShutdown());
+    assertFalse(executorService.isTerminated());
+
+    executorService.shutdownNow();
+    assertTrue(executorService.isShutdown());
+    assertTrue(executorService.isTerminated());
+    assertTrue(executorService.awaitTermination(
+        TestUtils.TEST_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS));
+    assertEquals(0, threadFactory.counter.get());
+  }
+
+  @Test
+  public void testSubmit() throws InterruptedException, ExecutionException {
+    CountingThreadFactory threadFactory = new CountingThreadFactory();
+    GaeExecutorService executorService = new GaeExecutorService("test", threadFactory);
+
+    final Semaphore semaphore = new Semaphore(0);
+    Future future = executorService.submit(new Runnable() {
+      @Override
+      public void run() {
+        semaphore.release();
+      }
+    });
+    assertNotNull(future);
+    waitFor(semaphore);
+
+    future = executorService.submit(new Runnable() {
+      @Override
+      public void run() {
+        semaphore.release();
+      }
+    }, "result");
+    assertNotNull(future);
+    waitFor(semaphore);
+    assertEquals("result", future.get());
+
+    future = executorService.submit(new Callable() {
+      @Override
+      public Object call() throws Exception {
+        semaphore.release();
+        return "result2";
+      }
+    });
+    assertNotNull(future);
+    waitFor(semaphore);
+    assertEquals("result2", future.get());
+
+    executorService.execute(new Runnable() {
+      @Override
+      public void run() {
+        semaphore.release();
+      }
+    });
+    waitFor(semaphore);
+
+    assertEquals(4, threadFactory.counter.get());
+
+    executorService.shutdown();
+    assertTrue(executorService.awaitTermination(
+        TestUtils.TEST_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS));
+    assertTrue(executorService.isShutdown());
+    assertTrue(executorService.isTerminated());
+  }
+
+  private static class CountingThreadFactory implements ThreadFactory {
+
+    private final AtomicInteger counter = new AtomicInteger(0);
+    private final ThreadFactory delegate = Executors.defaultThreadFactory();
+
+    @Override
+    public Thread newThread(Runnable r) {
+      counter.incrementAndGet();
+      return delegate.newThread(r);
+    }
+  }
+}