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));