MediaSession2: Add/remove playback listeners

Test: Run all MediaComponents test once
Change-Id: Ic24a67cbbead7a9d4d420fc03c8004cbd04f61b9
diff --git a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
index dba5e85..149765e 100644
--- a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
@@ -56,13 +56,6 @@
 
     private final MediaController2 mInstance;
 
-    /**
-     * Flag used by MediaController2Record to filter playback callback.
-     */
-    static final int CALLBACK_FLAG_PLAYBACK = 0x1;
-
-    static final int REQUEST_CODE_ALL = 0;
-
     private final Object mLock = new Object();
 
     private final Context mContext;
@@ -73,11 +66,11 @@
     private final IBinder.DeathRecipient mDeathRecipient;
 
     @GuardedBy("mLock")
-    private final List<PlaybackListenerHolder> mPlaybackListeners = new ArrayList<>();
-    @GuardedBy("mLock")
     private SessionServiceConnection mServiceConnection;
     @GuardedBy("mLock")
     private boolean mIsReleased;
+    @GuardedBy("mLock")
+    private PlaybackState2 mPlaybackState;
 
     // Assignment should be used with the lock hold, but should be used without a lock to prevent
     // potential deadlock.
@@ -185,7 +178,6 @@
                 mContext.unbindService(mServiceConnection);
                 mServiceConnection = null;
             }
-            mPlaybackListeners.clear();
             binder = mSessionBinder;
             mSessionBinder = null;
             mSessionCallbackStub.destroy();
@@ -373,8 +365,9 @@
 
     @Override
     public PlaybackState2 getPlaybackState_impl() {
-        // TODO(jaewan): Implement
-        return null;
+        synchronized (mLock) {
+            return mPlaybackState;
+        }
     }
 
     @Override
@@ -401,24 +394,15 @@
     ///////////////////////////////////////////////////
     // Protected or private methods
     ///////////////////////////////////////////////////
-    // Should be used without a lock to prevent potential deadlock.
-    private void registerCallbackForPlaybackNotLocked() {
-        final IMediaSession2 binder = mSessionBinder;
-        if (binder != null) {
-            try {
-                binder.registerCallback(mSessionCallbackStub,
-                        CALLBACK_FLAG_PLAYBACK, REQUEST_CODE_ALL);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Cannot connect to the service or the session is gone", e);
-            }
-        }
-    }
-
     private void pushPlaybackStateChanges(final PlaybackState2 state) {
         synchronized (mLock) {
-            for (int i = 0; i < mPlaybackListeners.size(); i++) {
-                mPlaybackListeners.get(i).postPlaybackChange(state);
-            }
+            mPlaybackState = state;
+            mCallbackExecutor.execute(() -> {
+                if (!mInstance.isConnected()) {
+                    return;
+                }
+                mCallback.onPlaybackStateChanged(state);
+            });
         }
     }
 
@@ -437,7 +421,6 @@
                 release = true;
                 return;
             }
-            boolean registerCallbackForPlaybackNeeded;
             synchronized (mLock) {
                 if (mIsReleased) {
                     return;
@@ -460,15 +443,11 @@
                     release = true;
                     return;
                 }
-                registerCallbackForPlaybackNeeded = !mPlaybackListeners.isEmpty();
             }
             // TODO(jaewan): Keep commands to prevents illegal API calls.
             mCallbackExecutor.execute(() -> {
                 mCallback.onConnected(commandGroup);
             });
-            if (registerCallbackForPlaybackNeeded) {
-                registerCallbackForPlaybackNotLocked();
-            }
         } finally {
             if (release) {
                 // Trick to call release() without holding the lock, to prevent potential deadlock
@@ -510,7 +489,13 @@
 
         @Override
         public void onPlaybackStateChanged(Bundle state) throws RuntimeException {
-            final MediaController2Impl controller = getController();
+            final MediaController2Impl controller;
+            try {
+                controller = getController();
+            } catch (IllegalStateException e) {
+                Log.w(TAG, "Don't fail silently here. Highly likely a bug");
+                return;
+            }
             controller.pushPlaybackStateChanges(PlaybackState2.fromBundle(state));
         }
 
diff --git a/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java b/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
index e173712..c2a3c5d 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
@@ -32,6 +32,7 @@
 import android.media.MediaItem2;
 import android.media.MediaLibraryService2;
 import android.media.MediaPlayerInterface;
+import android.media.MediaPlayerInterface.PlaybackListener;
 import android.media.MediaSession2;
 import android.media.MediaSession2.Builder;
 import android.media.MediaSession2.Command;
@@ -379,6 +380,44 @@
         mPlayer.setCurrentPlaylistItem(index);
     }
 
+    @Override
+    public void addPlaybackListener_impl(Executor executor, PlaybackListener listener) {
+        if (executor == null) {
+            throw new IllegalArgumentException("executor shouldn't be null");
+        }
+        if (listener == null) {
+            throw new IllegalArgumentException("listener shouldn't be null");
+        }
+        ensureCallingThread();
+        if (PlaybackListenerHolder.contains(mListeners, listener)) {
+            Log.w(TAG, "listener is already added. Ignoring.");
+            return;
+        }
+        mListeners.add(new PlaybackListenerHolder(executor, listener));
+        executor.execute(() -> listener.onPlaybackChanged(getInstance().getPlaybackState()));
+    }
+
+    @Override
+    public void removePlaybackListener_impl(PlaybackListener listener) {
+        if (listener == null) {
+            throw new IllegalArgumentException("listener shouldn't be null");
+        }
+        ensureCallingThread();
+        int idx = PlaybackListenerHolder.indexOf(mListeners, listener);
+        if (idx >= 0) {
+            mListeners.remove(idx);
+        }
+    }
+
+    @Override
+    public PlaybackState2 getPlaybackState_impl() {
+        ensureCallingThread();
+        ensurePlayer();
+        // TODO(jaewan): Is it safe to be called on any thread?
+        //               Otherwise we should cache the result from listener.
+        return mPlayer.getPlaybackState();
+    }
+
     ///////////////////////////////////////////////////
     // Protected or private methods
     ///////////////////////////////////////////////////
@@ -403,7 +442,6 @@
         }*/
     }
 
-
     private void ensurePlayer() {
         // TODO(jaewan): Should we pend command instead? Follow the decision from MP2.
         //               Alternatively we can add a API like setAcceptsPendingCommands(boolean).
@@ -473,11 +511,6 @@
         private final boolean mIsTrusted;
         private final IMediaSession2Callback mControllerBinder;
 
-        // Flag to indicate which callbacks should be returned for the controller binder.
-        // Either 0 or combination of {@link #CALLBACK_FLAG_PLAYBACK},
-        // {@link #CALLBACK_FLAG_SESSION_ACTIVENESS}
-        private int mFlag;
-
         public ControllerInfoImpl(Context context, ControllerInfo instance, int uid,
                 int pid, String packageName, IMediaSession2Callback callback) {
             mInstance = instance;
@@ -561,18 +594,6 @@
             return mControllerBinder;
         }
 
-        public boolean containsFlag(int flag) {
-            return (mFlag & flag) != 0;
-        }
-
-        public void addFlag(int flag) {
-            mFlag |= flag;
-        }
-
-        public void removeFlag(int flag) {
-            mFlag &= ~flag;
-        }
-
         public static ControllerInfoImpl from(ControllerInfo controller) {
             return (ControllerInfoImpl) controller.getProvider();
         }
diff --git a/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java b/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
index 77e7f22..a039d93 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
@@ -16,8 +16,6 @@
 
 package com.android.media;
 
-import static com.android.media.MediaController2Impl.CALLBACK_FLAG_PLAYBACK;
-
 import android.media.IMediaSession2;
 import android.media.IMediaSession2Callback;
 import android.media.MediaLibraryService2.BrowserRoot;
@@ -82,7 +80,7 @@
     }
 
     @Override
-    public void connect(String callingPackage, IMediaSession2Callback callback)
+    public void connect(String callingPackage, final IMediaSession2Callback callback)
             throws RuntimeException {
         final MediaSession2Impl sessionImpl = getSession();
         final ControllerInfo request = new ControllerInfo(sessionImpl.getContext(),
@@ -112,12 +110,29 @@
                         + " accept=" + accept);
             }
             try {
-                impl.getControllerBinder().onConnectionChanged(
+                callback.onConnectionChanged(
                         accept ? MediaSession2Stub.this : null,
                         allowedCommands == null ? null : allowedCommands.toBundle());
             } catch (RemoteException e) {
                 // Controller may be died prematurely.
             }
+            if (accept) {
+                // If connection is accepted, notify the current state to the controller.
+                // It's needed because we cannot call synchronous calls between session/controller.
+                // Note: We're doing this after the onConnectionChanged(), but there's no guarantee
+                //       that events here are notified after the onConnected() because
+                //       IMediaSession2Callback is oneway (i.e. async call) and CallbackStub will
+                //       use thread poll for incoming calls.
+                // TODO(jaewan): Should we protect getting playback state?
+                final PlaybackState2 state = session.getInstance().getPlaybackState();
+                final Bundle bundle = state != null ? state.toBundle() : null;
+                try {
+                    callback.onPlaybackStateChanged(bundle);
+                } catch (RemoteException e) {
+                    // TODO(jaewan): Handle this.
+                    // Controller may be died prematurely.
+                }
+            }
         });
     }
 
@@ -239,48 +254,13 @@
         });
     }
 
-    @Deprecated
-    @Override
-    public Bundle getPlaybackState() throws RemoteException {
-        MediaSession2Impl session = getSession();
-        // TODO(jaewan): Check if mPlayer.getPlaybackState() is safe here.
-        return session.getInstance().getPlayer().getPlaybackState().toBundle();
-    }
-
-    @Deprecated
-    @Override
-    public void registerCallback(final IMediaSession2Callback callbackBinder,
-            final int callbackFlag, final int requestCode) throws RemoteException {
-        // TODO(jaewan): Call onCommand() here. To do so, you should pend message.
-        synchronized (mLock) {
-            ControllerInfo controllerInfo = getController(callbackBinder);
-            if (controllerInfo == null) {
-                return;
-            }
-            ControllerInfoImpl.from(controllerInfo).addFlag(callbackFlag);
-        }
-    }
-
-    @Deprecated
-    @Override
-    public void unregisterCallback(IMediaSession2Callback callbackBinder, int callbackFlag)
-            throws RemoteException {
-        // TODO(jaewan): Call onCommand() here. To do so, you should pend message.
-        synchronized (mLock) {
-            ControllerInfo controllerInfo = getController(callbackBinder);
-            if (controllerInfo == null) {
-                return;
-            }
-            ControllerInfoImpl.from(controllerInfo).removeFlag(callbackFlag);
-        }
-    }
-
     private ControllerInfo getController(IMediaSession2Callback caller) {
         synchronized (mLock) {
             return mControllers.get(caller.asBinder());
         }
     }
 
+    // TODO(jaewan): Need a way to get controller with permissions
     public List<ControllerInfo> getControllers() {
         ArrayList<ControllerInfo> controllers = new ArrayList<>();
         synchronized (mLock) {
@@ -291,27 +271,15 @@
         return controllers;
     }
 
-    public List<ControllerInfo> getControllersWithFlag(int flag) {
-        ArrayList<ControllerInfo> controllers = new ArrayList<>();
-        synchronized (mLock) {
-            for (int i = 0; i < mControllers.size(); i++) {
-                ControllerInfo controllerInfo = mControllers.valueAt(i);
-                if (ControllerInfoImpl.from(controllerInfo).containsFlag(flag)) {
-                    controllers.add(controllerInfo);
-                }
-            }
-        }
-        return controllers;
-    }
-
     // Should be used without a lock to prevent potential deadlock.
     public void notifyPlaybackStateChangedNotLocked(PlaybackState2 state) {
-        final List<ControllerInfo> list = getControllersWithFlag(CALLBACK_FLAG_PLAYBACK);
+        final List<ControllerInfo> list = getControllers();
         for (int i = 0; i < list.size(); i++) {
             IMediaSession2Callback callbackBinder =
                     ControllerInfoImpl.from(list.get(i)).getControllerBinder();
             try {
-                callbackBinder.onPlaybackStateChanged(state.toBundle());
+                final Bundle bundle = state != null ? state.toBundle() : null;
+                callbackBinder.onPlaybackStateChanged(bundle);
             } catch (RemoteException e) {
                 Log.w(TAG, "Controller is gone", e);
                 // TODO(jaewan): What to do when the controller is gone?
diff --git a/packages/MediaComponents/test/runtest.sh b/packages/MediaComponents/test/runtest.sh
index d0290e7..920fa96 100644
--- a/packages/MediaComponents/test/runtest.sh
+++ b/packages/MediaComponents/test/runtest.sh
@@ -134,7 +134,7 @@
       ${adb} shell start
       ${adb} wait-for-device || break
       # Ensure package manager is loaded.
-      sleep 5
+      sleep 15
 
       # Install apks
       local install_failed="false"
diff --git a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
index ef75060..c1e9bdf 100644
--- a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
@@ -23,6 +23,7 @@
 import android.content.Context;
 import android.media.MediaBrowser2.BrowserCallback;
 import android.media.MediaSession2.CommandGroup;
+import android.media.MediaSession2.PlaylistParams;
 import android.os.Bundle;
 import android.support.annotation.CallSuper;
 import android.support.annotation.NonNull;
@@ -101,8 +102,26 @@
         }
 
         @Override
+        public void onPlaybackStateChanged(PlaybackState2 state) {
+            super.onPlaybackStateChanged(state);
+            if (mCallbackProxy != null) {
+                mCallbackProxy.onPlaybackStateChanged(state);
+            }
+        }
+
+        @Override
+        public void onPlaylistParamsChanged(PlaylistParams params) {
+            super.onPlaylistParamsChanged(params);
+            if (mCallbackProxy != null) {
+                mCallbackProxy.onPlaylistParamsChanged(params);
+            }
+        }
+
+        @Override
         public void onGetRootResult(Bundle rootHints, String rootMediaId, Bundle rootExtra) {
-            mCallbackProxy.onGetRootResult(rootHints, rootMediaId, rootExtra);
+            if (mCallbackProxy != null) {
+                mCallbackProxy.onGetRootResult(rootHints, rootMediaId, rootExtra);
+            }
         }
 
         @Override
diff --git a/packages/MediaComponents/test/src/android/media/MediaController2Test.java b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
index d7e0ae0..3b2afac 100644
--- a/packages/MediaComponents/test/src/android/media/MediaController2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
@@ -16,6 +16,7 @@
 
 package android.media;
 
+import android.media.MediaController2.ControllerCallback;
 import android.media.MediaPlayerInterface.PlaybackListener;
 import android.media.MediaSession2.ControllerInfo;
 import android.media.MediaSession2.SessionCallback;
@@ -198,62 +199,31 @@
         assertEquals(mContext.getPackageName(), mController.getSessionToken().getPackageName());
     }
 
+    // This also tests testGetPlaybackState().
     @Test
-    public void testGetPlaybackState() throws InterruptedException {
-        // TODO(jaewan): add equivalent test later
-        /*
-        final CountDownLatch latch = new CountDownLatch(1);
-        final MediaPlayerInterface.PlaybackListener listener = (state) -> {
-            assertEquals(PlaybackState.STATE_BUFFERING, state.getState());
-            latch.countDown();
-        };
-        assertNull(mController.getPlaybackState());
-        mController.addPlaybackListener(listener, sHandler);
-
-        mPlayer.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_BUFFERING));
-        assertTrue(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
-        assertEquals(PlaybackState.STATE_BUFFERING, mController.getPlaybackState().getState());
-        */
-    }
-
-    // TODO(jaewan): add equivalent test later
-    /*
-    @Test
-    public void testAddPlaybackListener() throws InterruptedException {
+    public void testControllerCallback_onPlaybackStateChanged() throws InterruptedException {
         final CountDownLatch latch = new CountDownLatch(2);
-        final MediaPlayerInterface.PlaybackListener listener = (state) -> {
-            switch ((int) latch.getCount()) {
-                case 2:
-                    assertEquals(PlaybackState.STATE_PLAYING, state.getState());
-                    break;
-                case 1:
-                    assertEquals(PlaybackState.STATE_PAUSED, state.getState());
-                    break;
+        final TestControllerCallbackInterface callback = new TestControllerCallbackInterface() {
+            @Override
+            public void onPlaybackStateChanged(PlaybackState2 state) {
+                switch ((int) latch.getCount()) {
+                    case 2:
+                        assertEquals(PlaybackState.STATE_PLAYING, state.getState());
+                        break;
+                    case 1:
+                        assertEquals(PlaybackState.STATE_PAUSED, state.getState());
+                        break;
+                }
+                latch.countDown();
             }
-            latch.countDown();
         };
 
-        mController.addPlaybackListener(listener, sHandler);
-        sHandler.postAndSync(()->{
-            mPlayer.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PLAYING));
-            mPlayer.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PAUSED));
-        });
-        assertTrue(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
-    }
-
-    @Test
-    public void testRemovePlaybackListener() throws InterruptedException {
-        final CountDownLatch latch = new CountDownLatch(1);
-        final MediaPlayerInterface.PlaybackListener listener = (state) -> {
-            fail();
-            latch.countDown();
-        };
-        mController.addPlaybackListener(listener, sHandler);
-        mController.removePlaybackListener(listener);
         mPlayer.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PLAYING));
-        assertFalse(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        mController = createController(mSession.getToken(), true, callback);
+        mPlayer.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PAUSED));
+        assertTrue(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
+        assertEquals(PlaybackState.STATE_PAUSED, mController.getPlaybackState().getState());
     }
-    */
 
     @Test
     public void testControllerCallback_onConnected() throws InterruptedException {
diff --git a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
index e7a0971..77893e9 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
@@ -19,11 +19,15 @@
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertTrue;
 
+import static android.media.TestUtils.createPlaybackState;
+
+import android.media.MediaPlayerInterface.PlaybackListener;
 import android.media.MediaSession2.Builder;
 import android.media.MediaSession2.ControllerInfo;
 import android.media.MediaSession2.PlaylistParams;
 import android.media.MediaSession2.SessionCallback;
 import android.os.Bundle;
+import android.os.Looper;
 import android.os.Process;
 import android.support.annotation.NonNull;
 import android.support.test.filters.SmallTest;
@@ -32,6 +36,7 @@
 import java.util.ArrayList;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -159,10 +164,10 @@
         assertTrue(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
     }
 
+    // TODO(jaewan): Re-enable test..
+    @Ignore
     @Test
     public void testPlaybackStateChangedListener() throws InterruptedException {
-        // TODO(jaewan): Add equivalent tests again
-        /*
         final CountDownLatch latch = new CountDownLatch(2);
         final MockPlayer player = new MockPlayer(0);
         final PlaybackListener listener = (state) -> {
@@ -170,45 +175,42 @@
             assertNotNull(state);
             switch ((int) latch.getCount()) {
                 case 2:
-                    assertEquals(PlaybackState.STATE_PLAYING, state.getState());
+                    assertEquals(PlaybackState2.STATE_PLAYING, state.getState());
                     break;
                 case 1:
-                    assertEquals(PlaybackState.STATE_PAUSED, state.getState());
+                    assertEquals(PlaybackState2.STATE_PAUSED, state.getState());
                     break;
                 case 0:
                     fail();
             }
             latch.countDown();
         };
-        player.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PLAYING));
+        player.notifyPlaybackState(createPlaybackState(PlaybackState2.STATE_PLAYING));
         sHandler.postAndSync(() -> {
-            mSession.addPlaybackListener(listener, sHandler);
+            mSession.addPlaybackListener(sHandlerExecutor, listener);
             // When the player is set, listeners will be notified about the player's current state.
             mSession.setPlayer(player);
         });
-        player.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PAUSED));
+        player.notifyPlaybackState(createPlaybackState(PlaybackState2.STATE_PAUSED));
         assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-        */
     }
 
     @Test
     public void testBadPlayer() throws InterruptedException {
         // TODO(jaewan): Add equivalent tests again
-        /*
         final CountDownLatch latch = new CountDownLatch(3); // expected call + 1
         final BadPlayer player = new BadPlayer(0);
         sHandler.postAndSync(() -> {
-            mSession.addPlaybackListener((state) -> {
+            mSession.addPlaybackListener(sHandlerExecutor, (state) -> {
                 // This will be called for every setPlayer() calls, but no more.
                 assertNull(state);
                 latch.countDown();
-            }, sHandler);
+            });
             mSession.setPlayer(player);
             mSession.setPlayer(mPlayer);
         });
-        player.notifyPlaybackState(createPlaybackState(PlaybackState.STATE_PAUSED));
+        player.notifyPlaybackState(createPlaybackState(PlaybackState2.STATE_PAUSED));
         assertFalse(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
-        */
     }
 
     private static class BadPlayer extends MockPlayer {
diff --git a/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java b/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
index 99ed4b9..d7c14e2 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
@@ -61,6 +61,9 @@
         // Add methods in ControllerCallback/BrowserCallback that you want to test.
         default void onPlaylistParamsChanged(MediaSession2.PlaylistParams params) {}
 
+        // Currently empty. Add methods in ControllerCallback/BrowserCallback that you want to test.
+        default void onPlaybackStateChanged(PlaybackState2 state) { }
+
         // Browser specific callbacks
         default void onGetRootResult(Bundle rootHints, String rootMediaId, Bundle rootExtra) {}
     }
@@ -177,6 +180,14 @@
         }
 
         @Override
+        public void onPlaybackStateChanged(PlaybackState2 state) {
+            super.onPlaybackStateChanged(state);
+            if (mCallbackProxy != null) {
+                mCallbackProxy.onPlaybackStateChanged(state);
+            }
+        }
+
+        @Override
         public void waitForConnect(boolean expect) throws InterruptedException {
             if (expect) {
                 assertTrue(connectLatch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));