System MR2: Fix bug on onTransfer and on getting initial routes

System MR2 should be able to get the routes right after creation,
if the app's router has already set the preference.

Bug: 181122201
Test: Test passed on transferTo.
Change-Id: I671d17b33ee43a939a02aee130333aef801d5789

2 files changed, 18 insertions, 1 deletions

diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index 5f44b62da6f7..9566e1c7aa01 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -260,6 +260,9 @@ public final class MediaRouter2 {
         mManagerCallback = new ManagerCallback();
         mHandler = new Handler(Looper.getMainLooper());
         mSystemController = new SystemRoutingController(sManager.getSystemRoutingSession());
+        mDiscoveryPreference = new RouteDiscoveryPreference.Builder(
+                sManager.getPreferredFeatures(clientPackageName), true).build();
+        updateAllRoutesFromManager();
         mMediaRouterService = null; // TODO: Make this non-null and check permission.
 
         // Only used by non-system MediaRouter2.
@@ -1946,7 +1949,7 @@ public final class MediaRouter2 {
             }
 
             RoutingController newController;
-            if (oldSession.isSystemSession()) {
+            if (newSession.isSystemSession()) {
                 mSystemController.setRoutingSessionInfo(newSession);
                 newController = mSystemController;
             } else {
diff --git a/media/java/android/media/MediaRouter2Manager.java b/media/java/android/media/MediaRouter2Manager.java
index 20e3573cfbe4..6fefbe15abae 100644
--- a/media/java/android/media/MediaRouter2Manager.java
+++ b/media/java/android/media/MediaRouter2Manager.java
@@ -237,6 +237,20 @@ public final class MediaRouter2Manager {
     }
 
     /**
+     * Returns the preferred features of the specified package name.
+     */
+    @NonNull
+    public List<String> getPreferredFeatures(@NonNull String packageName) {
+        Objects.requireNonNull(packageName, "packageName must not be null");
+
+        List<String> preferredFeatures = mPreferredFeaturesMap.get(packageName);
+        if (preferredFeatures == null) {
+            preferredFeatures = Collections.emptyList();
+        }
+        return preferredFeatures;
+    }
+
+    /**
      * Returns a list of routes which are related to the given package name in the given route list.
      */
     @NonNull