Merge "MediaSession2: Implement setAllowedCommands()" into pi-dev
diff --git a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
index 69d2fa8..e58bd02 100644
--- a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
@@ -339,7 +339,9 @@
         final CountDownLatch latch = new CountDownLatch(1);
         final SessionCallbackProxy callbackProxy = new SessionCallbackProxy(mContext) {
             @Override
-            public void onSubscribe(ControllerInfo info, String parentId, Bundle extras) {
+            public void onSubscribe(@NonNull MediaLibrarySession session,
+                    @NonNull ControllerInfo info, @NonNull String parentId,
+                    @Nullable Bundle extras) {
                 if (Process.myUid() == info.getUid()) {
                     assertEquals(testParentId, parentId);
                     assertTrue(TestUtils.equals(testExtras, extras));
@@ -361,7 +363,8 @@
         final CountDownLatch latch = new CountDownLatch(1);
         final SessionCallbackProxy callbackProxy = new SessionCallbackProxy(mContext) {
             @Override
-            public void onUnsubscribe(ControllerInfo info, String parentId) {
+            public void onUnsubscribe(@NonNull MediaLibrarySession session,
+                    @NonNull ControllerInfo info, @NonNull String parentId) {
                 if (Process.myUid() == info.getUid()) {
                     assertEquals(testParentId, parentId);
                     latch.countDown();
@@ -387,21 +390,29 @@
         final CountDownLatch latch = new CountDownLatch(3);
         final SessionCallbackProxy sessionCallbackProxy = new SessionCallbackProxy(mContext) {
             @Override
-            public CommandGroup onConnect(ControllerInfo controller) {
-                final MockMediaLibraryService2 service = (MockMediaLibraryService2)
-                        TestServiceRegistry.getInstance().getServiceInstance();
-                final MediaLibrarySession session = (MediaLibrarySession) service.getSession();
-                // Shouldn't trigger onChildrenChanged() for the browser, because it hasn't
-                // subscribed.
-                session.notifyChildrenChanged(testParentId1, testChildrenCount, null);
-                session.notifyChildrenChanged(controller, testParentId1, testChildrenCount, null);
-                return super.onConnect(controller);
+            public CommandGroup onConnect(@NonNull MediaSession2 session,
+                    @NonNull ControllerInfo controller) {
+                if (Process.myUid() == controller.getUid()) {
+                    assertTrue(session instanceof MediaLibrarySession);
+                    if (mSession != null) {
+                        mSession.close();
+                    }
+                    mSession = session;
+                    // Shouldn't trigger onChildrenChanged() for the browser, because it hasn't
+                    // subscribed.
+                    ((MediaLibrarySession) session).notifyChildrenChanged(
+                            testParentId1, testChildrenCount, null);
+                    ((MediaLibrarySession) session).notifyChildrenChanged(
+                            controller, testParentId1, testChildrenCount, null);
+                }
+                return super.onConnect(session, controller);
             }
 
             @Override
-            public void onSubscribe(ControllerInfo info, String parentId, Bundle extras) {
+            public void onSubscribe(@NonNull MediaLibrarySession session,
+                    @NonNull ControllerInfo info, @NonNull String parentId,
+                    @Nullable Bundle extras) {
                 if (Process.myUid() == info.getUid()) {
-                    final MediaLibrarySession session =  (MediaLibrarySession) mSession;
                     session.notifyChildrenChanged(testParentId2, testChildrenCount, null);
                     session.notifyChildrenChanged(info, testParentId2, testChildrenCount,
                             testExtras);
@@ -435,12 +446,6 @@
         final SessionToken2 token = MockMediaLibraryService2.getToken(mContext);
         final MediaBrowser2 browser = (MediaBrowser2) createController(
                 token, true, controllerCallbackProxy);
-        final MockMediaLibraryService2 service =
-                (MockMediaLibraryService2) TestServiceRegistry.getInstance().getServiceInstance();
-        if (mSession != null) {
-            mSession.close();
-        }
-        mSession = service.getSession();
         assertTrue(mSession instanceof MediaLibrarySession);
         browser.subscribe(testParentId2, null);
         // This ensures that onChildrenChanged() is only called for the expected reasons.
diff --git a/packages/MediaComponents/test/src/android/media/MediaController2Test.java b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
index 41f6aa4..bd5f031 100644
--- a/packages/MediaComponents/test/src/android/media/MediaController2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
@@ -32,6 +32,7 @@
 import android.os.HandlerThread;
 import android.os.Process;
 import android.os.ResultReceiver;
+import android.support.annotation.NonNull;
 import android.support.test.filters.FlakyTest;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
@@ -656,15 +657,6 @@
         assertEquals(SessionToken2.TYPE_SESSION_SERVICE, token.getType());
     }
 
-    private void connectToService(SessionToken2 token) throws InterruptedException {
-        if (mSession != null) {
-            mSession.close();
-        }
-        mController = createController(token);
-        mSession = TestServiceRegistry.getInstance().getServiceInstance().getSession();
-        mPlayer = (MockPlayer) mSession.getPlayer();
-    }
-
     @Test
     public void testConnectToService_sessionService() throws InterruptedException {
         testConnectToService(MockMediaSessionService2.ID);
@@ -680,17 +672,24 @@
         final CountDownLatch latch = new CountDownLatch(1);
         final SessionCallbackProxy proxy = new SessionCallbackProxy(mContext) {
             @Override
-            public CommandGroup onConnect(ControllerInfo controller) {
+            public CommandGroup onConnect(@NonNull MediaSession2 session,
+                    @NonNull ControllerInfo controller) {
                 if (Process.myUid() == controller.getUid()) {
+                    if (mSession != null) {
+                        mSession.close();
+                    }
+                    mSession = session;
+                    mPlayer = (MockPlayer) session.getPlayer();
                     assertEquals(mContext.getPackageName(), controller.getPackageName());
                     assertFalse(controller.isTrusted());
                     latch.countDown();
                 }
-                return super.onConnect(controller);
+                return super.onConnect(session, controller);
             }
         };
         TestServiceRegistry.getInstance().setSessionCallbackProxy(proxy);
-        connectToService(TestUtils.getServiceToken(mContext, id));
+
+        mController = createController(TestUtils.getServiceToken(mContext, id));
         assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
 
         // Test command from controller to session service
@@ -718,11 +717,14 @@
         testControllerAfterSessionIsGone(mSession.getToken().getId());
     }
 
+    // TODO(jaewan): Re-enable this test
     @Ignore
     @Test
     public void testControllerAfterSessionIsGone_sessionService() throws InterruptedException {
+        /*
         connectToService(TestUtils.getServiceToken(mContext, MockMediaSessionService2.ID));
         testControllerAfterSessionIsGone(MockMediaSessionService2.ID);
+        */
     }
 
     @Test
diff --git a/packages/MediaComponents/test/src/android/media/MockMediaLibraryService2.java b/packages/MediaComponents/test/src/android/media/MockMediaLibraryService2.java
index fb02f7a..df516c5 100644
--- a/packages/MediaComponents/test/src/android/media/MockMediaLibraryService2.java
+++ b/packages/MediaComponents/test/src/android/media/MockMediaLibraryService2.java
@@ -27,6 +27,8 @@
 import android.media.TestServiceRegistry.SessionCallbackProxy;
 import android.media.TestUtils.SyncHandler;
 import android.os.Bundle;
+import android.support.annotation.NonNull;
+import android.support.annotation.Nullable;
 import android.util.Log;
 
 import java.io.FileDescriptor;
@@ -96,8 +98,8 @@
 
     @Override
     public void onCreate() {
-        super.onCreate();
         TestServiceRegistry.getInstance().setServiceInstance(this);
+        super.onCreate();
     }
 
     @Override
@@ -144,9 +146,9 @@
         }
 
         @Override
-        public CommandGroup onConnect(MediaSession2 session,
-                ControllerInfo controller) {
-            return mCallbackProxy.onConnect(controller);
+        public CommandGroup onConnect(@NonNull MediaSession2 session,
+                @NonNull ControllerInfo controller) {
+            return mCallbackProxy.onConnect(session, controller);
         }
 
         @Override
@@ -211,15 +213,16 @@
         }
 
         @Override
-        public void onSubscribe(MediaLibrarySession session, ControllerInfo controller,
-                String parentId, Bundle extras) {
-            mCallbackProxy.onSubscribe(controller, parentId, extras);
+        public void onSubscribe(@NonNull MediaLibrarySession session,
+                @NonNull ControllerInfo controller, @NonNull String parentId,
+                @Nullable Bundle extras) {
+            mCallbackProxy.onSubscribe(session, controller, parentId, extras);
         }
 
         @Override
-        public void onUnsubscribe(MediaLibrarySession session, ControllerInfo controller,
-                String parentId) {
-            mCallbackProxy.onUnsubscribe(controller, parentId);
+        public void onUnsubscribe(@NonNull MediaLibrarySession session,
+                @NonNull ControllerInfo controller, String parentId) {
+            mCallbackProxy.onUnsubscribe(session, controller, parentId);
         }
     }
 
diff --git a/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java b/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
index ce7ce8b..64b1acd 100644
--- a/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
+++ b/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
@@ -27,6 +27,7 @@
 import android.media.MediaSession2.SessionCallback;
 import android.media.TestServiceRegistry.SessionCallbackProxy;
 import android.media.TestUtils.SyncHandler;
+import android.support.annotation.NonNull;
 
 import java.util.concurrent.Executor;
 
@@ -46,8 +47,8 @@
 
     @Override
     public void onCreate() {
-        super.onCreate();
         TestServiceRegistry.getInstance().setServiceInstance(this);
+        super.onCreate();
         mNotificationManager = (NotificationManager) getSystemService(Context.NOTIFICATION_SERVICE);
     }
 
@@ -103,9 +104,9 @@
         }
 
         @Override
-        public CommandGroup onConnect(MediaSession2 session,
-                ControllerInfo controller) {
-            return mCallbackProxy.onConnect(controller);
+        public CommandGroup onConnect(@NonNull MediaSession2 session,
+                @NonNull ControllerInfo controller) {
+            return mCallbackProxy.onConnect(session, controller);
         }
     }
 }
diff --git a/packages/MediaComponents/test/src/android/media/TestServiceRegistry.java b/packages/MediaComponents/test/src/android/media/TestServiceRegistry.java
index 08e0cf0..27b6f89 100644
--- a/packages/MediaComponents/test/src/android/media/TestServiceRegistry.java
+++ b/packages/MediaComponents/test/src/android/media/TestServiceRegistry.java
@@ -19,6 +19,7 @@
 import static org.junit.Assert.fail;
 
 import android.content.Context;
+import android.media.MediaLibraryService2.MediaLibrarySession;
 import android.media.MediaSession2.CommandGroup;
 import android.media.MediaSession2.ControllerInfo;
 import android.media.TestUtils.SyncHandler;
@@ -26,6 +27,8 @@
 import android.os.Handler;
 import android.os.Process;
 import android.support.annotation.GuardedBy;
+import android.support.annotation.NonNull;
+import android.support.annotation.Nullable;
 
 /**
  * Keeps the instance of currently running {@link MockMediaSessionService2}. And also provides
@@ -36,7 +39,7 @@
 public class TestServiceRegistry {
     /**
      * Proxy for both {@link MediaSession2.SessionCallback} and
-     * {@link MediaLibraryService2.MediaLibrarySessionCallback}.
+     * {@link MediaLibraryService2.MediaLibrarySession.MediaLibrarySessionCallback}.
      */
     public static abstract class SessionCallbackProxy {
         private final Context mContext;
@@ -56,7 +59,8 @@
          * @param controller
          * @return
          */
-        public CommandGroup onConnect(ControllerInfo controller) {
+        public CommandGroup onConnect(@NonNull MediaSession2 session,
+                @NonNull ControllerInfo controller) {
             if (Process.myUid() == controller.getUid()) {
                 CommandGroup commands = new CommandGroup(mContext);
                 commands.addAllPredefinedCommands();
@@ -75,8 +79,10 @@
          */
         public void onServiceDestroyed() { }
 
-        public void onSubscribe(ControllerInfo info, String parentId, Bundle extra) { }
-        public void onUnsubscribe(ControllerInfo info, String parentId) { }
+        public void onSubscribe(@NonNull MediaLibrarySession session, @NonNull ControllerInfo info,
+                @NonNull String parentId, @Nullable Bundle extra) { }
+        public void onUnsubscribe(@NonNull MediaLibrarySession session,
+                @NonNull ControllerInfo info, @NonNull String parentId) { }
     }
 
     @GuardedBy("TestServiceRegistry.class")