Merge "MediaSession2: Fix bug in MediaSession2.setCustomLayout()"
diff --git a/packages/MediaComponents/src/com/android/media/MediaBrowser2Impl.java b/packages/MediaComponents/src/com/android/media/MediaBrowser2Impl.java
index 3e6d98f..7e928d7 100644
--- a/packages/MediaComponents/src/com/android/media/MediaBrowser2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaBrowser2Impl.java
@@ -91,10 +91,4 @@
             mCallback.onGetRootResult(rootHints, rootMediaId, rootExtra);
         });
     }
-
-    public void onCustomLayoutChanged(final List<CommandButton> layout) {
-        getCallbackExecutor().execute(() -> {
-            mCallback.onCustomLayoutChanged(layout);
-        });
-    }
 }
diff --git a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
index 3d67eaa..5ae37ee 100644
--- a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
@@ -26,6 +26,7 @@
 import android.media.MediaItem2;
 import android.media.MediaSession2;
 import android.media.MediaSession2.Command;
+import android.media.MediaSession2.CommandButton;
 import android.media.MediaSession2.CommandGroup;
 import android.media.MediaController2;
 import android.media.MediaController2.ControllerCallback;
@@ -632,6 +633,12 @@
         });
     }
 
+    void onCustomLayoutChanged(final List<CommandButton> layout) {
+        mCallbackExecutor.execute(() -> {
+            mCallback.onCustomLayoutChanged(layout);
+        });
+    }
+
     // This will be called on the main thread.
     private class SessionServiceConnection implements ServiceConnection {
         @Override
diff --git a/packages/MediaComponents/src/com/android/media/MediaSession2CallbackStub.java b/packages/MediaComponents/src/com/android/media/MediaSession2CallbackStub.java
index bf88283..07edf7e 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSession2CallbackStub.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSession2CallbackStub.java
@@ -53,7 +53,6 @@
         return controller;
     }
 
-    // TODO(jaewan): Refactor code to get rid of these pattern.
     private MediaBrowser2Impl getBrowser() throws IllegalStateException {
         final MediaController2Impl controller = getController();
         if (controller instanceof MediaBrowser2Impl) {
@@ -163,49 +162,31 @@
     }
 
     @Override
-    public void onGetRootResult(Bundle rootHints, String rootMediaId, Bundle rootExtra)
-            throws RuntimeException {
-        final MediaBrowser2Impl browser;
-        try {
-            browser = getBrowser();
-        } catch (IllegalStateException e) {
-            Log.w(TAG, "Don't fail silently here. Highly likely a bug");
-            return;
-        }
-        if (browser == null) {
-            // TODO(jaewan): Revisit here. Could be a bug
-            return;
-        }
-        browser.onGetRootResult(rootHints, rootMediaId, rootExtra);
-    }
-
-    @Override
     public void onCustomLayoutChanged(List<Bundle> commandButtonlist) {
         if (commandButtonlist == null) {
             // Illegal call. Ignore
             return;
         }
-        // TODO(jaewan): Fix here. It's controller feature so shouldn't use browser
-        final MediaBrowser2Impl browser;
+        final MediaController2Impl controller;
         try {
-            browser = getBrowser();
+            controller = getController();
         } catch (IllegalStateException e) {
             Log.w(TAG, "Don't fail silently here. Highly likely a bug");
             return;
         }
-        if (browser == null) {
+        if (controller == null) {
             // TODO(jaewan): Revisit here. Could be a bug
             return;
         }
         List<CommandButton> layout = new ArrayList<>();
         for (int i = 0; i < commandButtonlist.size(); i++) {
             CommandButton button = CommandButtonImpl.fromBundle(
-                    browser.getContext(), commandButtonlist.get(i));
+                    controller.getContext(), commandButtonlist.get(i));
             if (button != null) {
                 layout.add(button);
             }
         }
-        browser.onCustomLayoutChanged(layout);
+        controller.onCustomLayoutChanged(layout);
     }
 
     @Override
@@ -223,4 +204,24 @@
         }
         controller.onCustomCommand(command, args, receiver);
     }
+
+    ////////////////////////////////////////////////////////////////////////////////////////////
+    // MediaBrowser specific
+    ////////////////////////////////////////////////////////////////////////////////////////////
+    @Override
+    public void onGetRootResult(Bundle rootHints, String rootMediaId, Bundle rootExtra)
+            throws RuntimeException {
+        final MediaBrowser2Impl browser;
+        try {
+            browser = getBrowser();
+        } catch (IllegalStateException e) {
+            Log.w(TAG, "Don't fail silently here. Highly likely a bug");
+            return;
+        }
+        if (browser == null) {
+            // TODO(jaewan): Revisit here. Could be a bug
+            return;
+        }
+        browser.onGetRootResult(rootHints, rootMediaId, rootExtra);
+    }
 }
diff --git a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
index 4cd8177..eff4c3b 100644
--- a/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaBrowser2Test.java
@@ -24,6 +24,7 @@
 import android.content.Context;
 import android.media.MediaBrowser2.BrowserCallback;
 import android.media.MediaSession2.Command;
+import android.media.MediaSession2.CommandButton;
 import android.media.MediaSession2.CommandGroup;
 import android.media.MediaSession2.PlaylistParams;
 import android.os.Bundle;
@@ -36,6 +37,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -104,42 +106,41 @@
         @CallSuper
         @Override
         public void onConnected(CommandGroup commands) {
-            super.onConnected(commands);
             connectLatch.countDown();
         }
 
         @CallSuper
         @Override
         public void onDisconnected() {
-            super.onDisconnected();
             disconnectLatch.countDown();
         }
 
         @Override
         public void onPlaybackStateChanged(PlaybackState2 state) {
-            super.onPlaybackStateChanged(state);
             mCallbackProxy.onPlaybackStateChanged(state);
         }
 
         @Override
         public void onPlaylistParamsChanged(PlaylistParams params) {
-            super.onPlaylistParamsChanged(params);
             mCallbackProxy.onPlaylistParamsChanged(params);
         }
 
         @Override
         public void onPlaybackInfoChanged(MediaController2.PlaybackInfo info) {
-            if (mCallbackProxy != null) {
-                mCallbackProxy.onPlaybackInfoChanged(info);
-            }
+            mCallbackProxy.onPlaybackInfoChanged(info);
         }
 
         @Override
         public void onCustomCommand(Command command, Bundle args, ResultReceiver receiver) {
-            super.onCustomCommand(command, args, receiver);
             mCallbackProxy.onCustomCommand(command, args, receiver);
         }
 
+
+        @Override
+        public void onCustomLayoutChanged(List<CommandButton> layout) {
+            mCallbackProxy.onCustomLayoutChanged(layout);
+        }
+
         @Override
         public void onGetRootResult(Bundle rootHints, String rootMediaId, Bundle rootExtra) {
             super.onGetRootResult(rootHints, rootMediaId, rootExtra);
diff --git a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
index b00633b..f5ac6aa 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
@@ -33,6 +33,8 @@
 import android.media.MediaPlayerInterface.PlaybackListener;
 import android.media.MediaSession2.Builder;
 import android.media.MediaSession2.Command;
+import android.media.MediaSession2.CommandButton;
+import android.media.MediaSession2.CommandGroup;
 import android.media.MediaSession2.ControllerInfo;
 import android.media.MediaSession2.PlaylistParams;
 import android.media.MediaSession2.SessionCallback;
@@ -353,6 +355,45 @@
     }
 
     @Test
+    public void testSetCustomLayout() throws InterruptedException {
+        final List<CommandButton> buttons = new ArrayList<>();
+        buttons.add(new CommandButton.Builder(mContext)
+                .setCommand(new Command(mContext, MediaSession2.COMMAND_CODE_PLAYBACK_PLAY))
+                .setDisplayName("button").build());
+        final CountDownLatch latch = new CountDownLatch(1);
+        final SessionCallback sessionCallback = new SessionCallback(mContext) {
+            @Override
+            public CommandGroup onConnect(ControllerInfo controller) {
+                if (mContext.getPackageName().equals(controller.getPackageName())) {
+                    mSession.setCustomLayout(controller, buttons);
+                }
+                return super.onConnect(controller);
+            }
+        };
+
+        try (final MediaSession2 session = new MediaSession2.Builder(mContext, mPlayer)
+                .setId("testSetCustomLayout")
+                .setSessionCallback(sHandlerExecutor, sessionCallback)
+                .build()) {
+            final TestControllerCallbackInterface callback = new TestControllerCallbackInterface() {
+                @Override
+                public void onCustomLayoutChanged(List<CommandButton> layout) {
+                    assertEquals(layout.size(), buttons.size());
+                    for (int i = 0; i < layout.size(); i++) {
+                        assertEquals(layout.get(i).getCommand(), buttons.get(i).getCommand());
+                        assertEquals(layout.get(i).getDisplayName(),
+                                buttons.get(i).getDisplayName());
+                    }
+                    latch.countDown();
+                }
+            };
+            final MediaController2 controller =
+                    createController(session.getToken(), true, callback);
+            assertTrue(latch.await(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
+        }
+    }
+
+    @Test
     public void testSendCustomAction() throws InterruptedException {
         final Command testCommand =
                 new Command(mContext, MediaSession2.COMMAND_CODE_PLAYBACK_PREPARE);
diff --git a/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java b/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
index 513fa29..f5abfff 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSession2TestBase.java
@@ -22,6 +22,7 @@
 import android.content.Context;
 import android.media.MediaController2.ControllerCallback;
 import android.media.MediaSession2.Command;
+import android.media.MediaSession2.CommandButton;
 import android.media.MediaSession2.CommandGroup;
 import android.os.Bundle;
 import android.os.HandlerThread;
@@ -59,15 +60,15 @@
         ControllerCallback getCallback();
     }
 
+    // Any change here should be also reflected to the TestControllerCallback and
+    // TestBrowserCallback
     interface TestControllerCallbackInterface {
-        // Add methods in ControllerCallback/BrowserCallback that you want to test.
+        // Add methods in ControllerCallback that you want to test.
         default void onPlaylistChanged(List<MediaItem2> playlist) {}
         default void onPlaylistParamsChanged(MediaSession2.PlaylistParams params) {}
         default void onPlaybackInfoChanged(MediaController2.PlaybackInfo info) {}
-
-        // Currently empty. Add methods in ControllerCallback/BrowserCallback that you want to test.
-        default void onPlaybackStateChanged(PlaybackState2 state) { }
-
+        default void onPlaybackStateChanged(PlaybackState2 state) {}
+        default void onCustomLayoutChanged(List<CommandButton> layout) {}
         default void onCustomCommand(Command command, Bundle args, ResultReceiver receiver) {}
     }
 
@@ -188,26 +189,22 @@
         @CallSuper
         @Override
         public void onConnected(CommandGroup commands) {
-            super.onConnected(commands);
             connectLatch.countDown();
         }
 
         @CallSuper
         @Override
         public void onDisconnected() {
-            super.onDisconnected();
             disconnectLatch.countDown();
         }
 
         @Override
         public void onPlaybackStateChanged(PlaybackState2 state) {
-            super.onPlaybackStateChanged(state);
             mCallbackProxy.onPlaybackStateChanged(state);
         }
 
         @Override
         public void onCustomCommand(Command command, Bundle args, ResultReceiver receiver) {
-            super.onCustomCommand(command, args, receiver);
             mCallbackProxy.onCustomCommand(command, args, receiver);
         }
 
@@ -231,23 +228,22 @@
 
         @Override
         public void onPlaylistChanged(List<MediaItem2> params) {
-            if (mCallbackProxy != null) {
-                mCallbackProxy.onPlaylistChanged(params);
-            }
+            mCallbackProxy.onPlaylistChanged(params);
         }
 
         @Override
         public void onPlaylistParamsChanged(MediaSession2.PlaylistParams params) {
-            if (mCallbackProxy != null) {
-                mCallbackProxy.onPlaylistParamsChanged(params);
-            }
+            mCallbackProxy.onPlaylistParamsChanged(params);
         }
 
         @Override
         public void onPlaybackInfoChanged(MediaController2.PlaybackInfo info) {
-            if (mCallbackProxy != null) {
-                mCallbackProxy.onPlaybackInfoChanged(info);
-            }
+            mCallbackProxy.onPlaybackInfoChanged(info);
+        }
+
+        @Override
+        public void onCustomLayoutChanged(List<CommandButton> layout) {
+            mCallbackProxy.onCustomLayoutChanged(layout);
         }
     }