Media: Remove @MainThread from MediaRouter2

Now it uses mLock to synchronize mCallbackRecords, mControlCategories,
and mClient instead of enforcing main thread.

Also, check concurrency of MediaRouter2 and MediaRouter2Manager.

This CL also removes control category parameter from addCallback(registerCallback)
not to permit two callbacks in an application to have different control categories.

Function names are changed from add/remove callbacks to register/unregister callbacks
to follow the API guideline.

Test: atest mediaroutertest

Change-Id: I40f2bc45f7d108416182b68b4ae23127f8bb2eb4

3 files changed, 201 insertions, 176 deletions

diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index a69b105f0f50..f4dffa28f133 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -16,12 +16,10 @@
 
 package android.media;
 
-import android.annotation.MainThread;
+import android.annotation.CallbackExecutor;
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.content.Context;
 import android.content.Intent;
-import android.os.Looper;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.UserHandle;
@@ -30,9 +28,11 @@ import android.util.Log;
 import com.android.internal.annotations.GuardedBy;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Executor;
 
 
@@ -43,7 +43,6 @@ import java.util.concurrent.Executor;
 public class MediaRouter2 {
     private static final String TAG = "MediaRouter";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
-
     private static final Object sLock = new Object();
 
     @GuardedBy("sLock")
@@ -51,134 +50,160 @@ public class MediaRouter2 {
 
     private Context mContext;
     private final IMediaRouterService mMediaRouterService;
-    private List<CallbackRecord> mCallbackRecords = new ArrayList<>();
-    final String mPackageName;
 
-    IMediaRouter2Client mClient;
+    private CopyOnWriteArrayList<CallbackRecord> mCallbackRecords = new CopyOnWriteArrayList<>();
+    @GuardedBy("sLock")
+    private List<String> mControlCategories = Collections.emptyList();
+    @GuardedBy("sLock")
+    private Client mClient;
+
+    private final String mPackageName;
 
     /**
      * Gets an instance of the media router associated with the context.
      */
     public static MediaRouter2 getInstance(@NonNull Context context) {
+        Objects.requireNonNull(context, "context must not be null");
         synchronized (sLock) {
             if (sInstance == null) {
-                sInstance = new MediaRouter2(context);
+                sInstance = new MediaRouter2(context.getApplicationContext());
             }
             return sInstance;
         }
     }
 
-    private MediaRouter2(Context context) {
-        mContext = Objects.requireNonNull(context, "context must not be null");
+    private MediaRouter2(Context appContext) {
+        mContext = appContext;
         mMediaRouterService = IMediaRouterService.Stub.asInterface(
                 ServiceManager.getService(Context.MEDIA_ROUTER_SERVICE));
         mPackageName = mContext.getPackageName();
+        //TODO: read control categories from the manifest
     }
 
     /**
-     * Registers a callback to discover routes that match the selector and to receive events
-     * when they change.
+     * Registers a callback to discover routes and to receive events when they change.
      */
-    @MainThread
-    public void addCallback(@NonNull List<String> controlCategories, @NonNull Executor executor,
+    public void registerCallback(@NonNull @CallbackExecutor Executor executor,
             @NonNull Callback callback) {
-        addCallback(controlCategories, executor, callback, 0);
+        registerCallback(executor, callback, 0);
     }
 
     /**
-     * Registers a callback to discover routes that match the selector and to receive events
-     * when they change.
+     * Registers a callback to discover routes and to receive events when they change.
      * <p>
-     * If you add the same callback twice or more, the previous arguments will be overwritten
+     * If you register the same callback twice or more, the previous arguments will be overwritten
      * with the new arguments.
      * </p>
      */
-    @MainThread
-    public void addCallback(@NonNull List<String> controlCategories, @NonNull Executor executor,
+    public void registerCallback(@NonNull @CallbackExecutor Executor executor,
             @NonNull Callback callback, int flags) {
-        checkMainThread();
-
-        Objects.requireNonNull(controlCategories, "control categories must not be null");
         Objects.requireNonNull(executor, "executor must not be null");
         Objects.requireNonNull(callback, "callback must not be null");
 
-        if (mCallbackRecords.size() == 0) {
-            Client client = new Client();
-            try {
-                mMediaRouterService.registerClient2AsUser(client, mPackageName,
-                        UserHandle.myUserId());
-                //TODO: We should merge control categories of callbacks.
-                mMediaRouterService.setControlCategories(client, controlCategories);
-                mClient = client;
-            } catch (RemoteException ex) {
-                Log.e(TAG, "Unable to register media router.", ex);
+        // This is required to prevent adding the same callback twice.
+        synchronized (mCallbackRecords) {
+            if (mCallbackRecords.size() == 0) {
+                synchronized (sLock) {
+                    Client client = new Client();
+                    try {
+                        mMediaRouterService.registerClient2AsUser(client, mPackageName,
+                                UserHandle.myUserId());
+                        //TODO: We should merge control categories of callbacks.
+                        mMediaRouterService.setControlCategories(client, mControlCategories);
+                        mClient = client;
+                    } catch (RemoteException ex) {
+                        Log.e(TAG, "Unable to register media router.", ex);
+                    }
+                }
             }
-        }
 
-        final int index = findCallbackRecordIndex(callback);
-        CallbackRecord record;
-        if (index < 0) {
-            record = new CallbackRecord(callback);
-            mCallbackRecords.add(record);
-        } else {
-            record = mCallbackRecords.get(index);
+            final int index = findCallbackRecordIndexLocked(callback);
+            CallbackRecord record;
+            if (index < 0) {
+                record = new CallbackRecord(callback);
+                mCallbackRecords.add(record);
+            } else {
+                record = mCallbackRecords.get(index);
+            }
+            record.mExecutor = executor;
+            record.mFlags = flags;
         }
-        record.mExecutor = executor;
-        record.mControlCategories = controlCategories;
-        record.mFlags = flags;
 
-        //TODO: Check if we need an update.
+        //TODO: Update discovery request here.
     }
 
     /**
-     * Removes the given callback. The callback will no longer receive events.
+     * Unregisters the given callback. The callback will no longer receive events.
      * If the callback has not been added or been removed already, it is ignored.
-     * @param callback the callback to remove.
-     * @see #addCallback
+     *
+     * @param callback the callback to unregister
+     * @see #registerCallback
      */
-    @MainThread
-    public void removeCallback(@NonNull Callback callback) {
-        checkMainThread();
-
+    public void unregisterCallback(@NonNull Callback callback) {
         Objects.requireNonNull(callback, "callback must not be null");
 
-        final int index = findCallbackRecordIndex(callback);
-        if (index < 0) {
-            Log.w(TAG, "Ignoring to remove unknown callback. " + callback);
-            return;
-        }
-        mCallbackRecords.remove(index);
-        if (mCallbackRecords.size() == 0 && mClient != null) {
-            try {
-                mMediaRouterService.unregisterClient2(mClient);
-            } catch (RemoteException ex) {
-                Log.e(TAG, "Unable to unregister media router.", ex);
+        synchronized (mCallbackRecords) {
+            final int index = findCallbackRecordIndexLocked(callback);
+            if (index < 0) {
+                Log.w(TAG, "Ignoring to remove unknown callback. " + callback);
+                return;
+            }
+            mCallbackRecords.remove(index);
+            synchronized (sLock) {
+                if (mCallbackRecords.size() == 0 && mClient != null) {
+                    try {
+                        mMediaRouterService.unregisterClient2(mClient);
+                    } catch (RemoteException ex) {
+                        Log.e(TAG, "Unable to unregister media router.", ex);
+                    }
+                    mClient = null;
+                }
             }
-            mClient = null;
         }
     }
 
-    private int findCallbackRecordIndex(Callback callback) {
-        final int count = mCallbackRecords.size();
-        for (int i = 0; i < count; i++) {
-            CallbackRecord callbackRecord = mCallbackRecords.get(i);
-            if (callbackRecord.mCallback == callback) {
-                return i;
+    //TODO(b/139033746): Rename "Control Category" when it's finalized.
+    /**
+     * Sets the control categories of the application.
+     * Routes that support at least one of the given control categories only exists and are handled
+     * by the media router.
+     */
+    public void setControlCategories(@NonNull Collection<String> controlCategories) {
+        Objects.requireNonNull(controlCategories, "control categories must not be null");
+
+        Client client;
+        List<String> newControlCategories;
+        synchronized (sLock) {
+            mControlCategories = new ArrayList<>(controlCategories);
+            newControlCategories = mControlCategories;
+            client = mClient;
+        }
+        if (client != null) {
+            try {
+                mMediaRouterService.setControlCategories(client, newControlCategories);
+            } catch (RemoteException ex) {
+                Log.e(TAG, "Unable to set control categories.", ex);
             }
         }
-        return -1;
     }
 
+
     /**
      * Selects the specified route.
      *
-     * @param route The route to select.
+     * @param route the route to select
      */
     //TODO: add a parameter for category (e.g. mirroring/casting)
-    public void selectRoute(@Nullable MediaRoute2Info route) {
-        if (mClient != null) {
+    public void selectRoute(@NonNull MediaRoute2Info route) {
+        Objects.requireNonNull(route, "route must not be null");
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
             try {
-                mMediaRouterService.selectRoute2(mClient, route);
+                mMediaRouterService.selectRoute2(client, route);
             } catch (RemoteException ex) {
                 Log.e(TAG, "Unable to select route.", ex);
             }
@@ -187,6 +212,7 @@ public class MediaRouter2 {
 
     /**
      * Sends a media control request to be performed asynchronously by the route's destination.
+     *
      * @param route the route that will receive the control request
      * @param request the media control request
      */
@@ -196,20 +222,29 @@ public class MediaRouter2 {
         Objects.requireNonNull(route, "route must not be null");
         Objects.requireNonNull(request, "request must not be null");
 
-        if (mClient != null) {
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
             try {
-                mMediaRouterService.sendControlRequest(mClient, route, request);
+                mMediaRouterService.sendControlRequest(client, route, request);
             } catch (RemoteException ex) {
                 Log.e(TAG, "Unable to send control request.", ex);
             }
         }
     }
 
-    void checkMainThread() {
-        Looper looper = Looper.myLooper();
-        if (looper == null || looper != Looper.getMainLooper()) {
-            throw new IllegalStateException("the method must be called on the main thread");
+    @GuardedBy("mCallbackRecords")
+    private int findCallbackRecordIndexLocked(Callback callback) {
+        final int count = mCallbackRecords.size();
+        for (int i = 0; i < count; i++) {
+            CallbackRecord callbackRecord = mCallbackRecords.get(i);
+            if (callbackRecord.mCallback == callback) {
+                return i;
+            }
         }
+        return -1;
     }
 
     /**
@@ -232,15 +267,13 @@ public class MediaRouter2 {
         public void onRouteRemoved(MediaRoute2Info routeInfo) {}
     }
 
-    final class CallbackRecord {
+    static final class CallbackRecord {
         public final Callback mCallback;
         public Executor mExecutor;
-        public List<String> mControlCategories;
         public int mFlags;
 
         CallbackRecord(@NonNull Callback callback) {
-            mCallback = Objects.requireNonNull(callback, "callback must not be null");
-            mControlCategories = Collections.emptyList();
+            mCallback = callback;
         }
     }
 
diff --git a/media/java/android/media/MediaRouter2Manager.java b/media/java/android/media/MediaRouter2Manager.java
index 85105e069aaa..6c53f7d6d73f 100644
--- a/media/java/android/media/MediaRouter2Manager.java
+++ b/media/java/android/media/MediaRouter2Manager.java
@@ -56,11 +56,10 @@ public class MediaRouter2Manager {
     final String mPackageName;
 
     private Context mContext;
+    @GuardedBy("sLock")
     private Client mClient;
     private final IMediaRouterService mMediaRouterService;
     final Handler mHandler;
-
-    @GuardedBy("sLock")
     final List<CallbackRecord> mCallbacks = new CopyOnWriteArrayList<>();
 
     @SuppressWarnings("WeakerAccess") /* synthetic access */
@@ -73,6 +72,7 @@ public class MediaRouter2Manager {
 
     /**
      * Gets an instance of media router manager that controls media route of other applications.
+     *
      * @return The media router manager instance for the context.
      */
     public static MediaRouter2Manager getInstance(@NonNull Context context) {
@@ -96,28 +96,29 @@ public class MediaRouter2Manager {
     /**
      * Registers a callback to listen route info.
      *
-     * @param executor The executor that runs the callback.
-     * @param callback The callback to add.
+     * @param executor the executor that runs the callback
+     * @param callback the callback to add
      */
-    public void addCallback(@NonNull @CallbackExecutor Executor executor,
+    public void registerCallback(@NonNull @CallbackExecutor Executor executor,
             @NonNull Callback callback) {
-
         Objects.requireNonNull(executor, "executor must not be null");
         Objects.requireNonNull(callback, "callback must not be null");
 
-        synchronized (sLock) {
-            if (findCallbackRecordIndex(callback) >= 0) {
+        synchronized (mCallbacks) {
+            if (findCallbackRecordIndexLocked(callback) >= 0) {
                 Log.w(TAG, "Ignoring to add the same callback twice.");
                 return;
             }
-            if (mCallbacks.size() == 0) {
-                Client client = new Client();
-                try {
-                    mMediaRouterService.registerManagerAsUser(client, mPackageName,
-                            UserHandle.myUserId());
-                    mClient = client;
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to register media router manager.", ex);
+            synchronized (sLock) {
+                if (mCallbacks.size() == 0) {
+                    Client client = new Client();
+                    try {
+                        mMediaRouterService.registerManagerAsUser(client, mPackageName,
+                                UserHandle.myUserId());
+                        mClient = client;
+                    } catch (RemoteException ex) {
+                        Log.e(TAG, "Unable to register media router manager.", ex);
+                    }
                 }
             }
             CallbackRecord record = new CallbackRecord(executor, callback);
@@ -127,34 +128,35 @@ public class MediaRouter2Manager {
     }
 
     /**
-     * Removes the specified callback.
+     * Unregisters the specified callback.
      *
-     * @param callback The callback to remove.
+     * @param callback the callback to unregister
      */
-    public void removeCallback(@NonNull Callback callback) {
-        if (callback == null) {
-            throw new IllegalArgumentException("callback must not be null");
-        }
+    public void unregisterCallback(@NonNull Callback callback) {
+        Objects.requireNonNull(callback, "callback must not be null");
 
-        synchronized (sLock) {
-            final int index = findCallbackRecordIndex(callback);
+        synchronized (mCallbacks) {
+            final int index = findCallbackRecordIndexLocked(callback);
             if (index < 0) {
                 Log.w(TAG, "Ignore removing unknown callback. " + callback);
                 return;
             }
             mCallbacks.remove(index);
-            if (mCallbacks.size() == 0 && mClient != null) {
-                try {
-                    mMediaRouterService.unregisterManager(mClient);
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to unregister media router manager", ex);
+            synchronized (sLock) {
+                if (mCallbacks.size() == 0 && mClient != null) {
+                    try {
+                        mMediaRouterService.unregisterManager(mClient);
+                    } catch (RemoteException ex) {
+                        Log.e(TAG, "Unable to unregister media router manager", ex);
+                    }
+                    mClient = null;
                 }
-                mClient = null;
             }
         }
     }
 
-    private int findCallbackRecordIndex(Callback callback) {
+    @GuardedBy("mCallbacks")
+    private int findCallbackRecordIndexLocked(Callback callback) {
         final int count = mCallbacks.size();
         for (int i = 0; i < count; i++) {
             if (mCallbacks.get(i).mCallback == callback) {
@@ -173,6 +175,8 @@ public class MediaRouter2Manager {
      */
     @NonNull
     public List<MediaRoute2Info> getAvailableRoutes(@NonNull String packageName) {
+        Objects.requireNonNull(packageName, "packageName must not be null");
+
         List<String> controlCategories = mControlCategoryMap.get(packageName);
         if (controlCategories == null) {
             return Collections.emptyList();
@@ -193,9 +197,16 @@ public class MediaRouter2Manager {
      * @param route the route to be selected
      */
     public void selectRoute(@NonNull String packageName, @NonNull MediaRoute2Info route) {
-        if (mClient != null) {
+        Objects.requireNonNull(packageName, "packageName must not be null");
+        Objects.requireNonNull(route, "route must not be null");
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
             try {
-                mMediaRouterService.selectClientRoute2(mClient, packageName, route);
+                mMediaRouterService.selectClientRoute2(client, packageName, route);
             } catch (RemoteException ex) {
                 Log.e(TAG, "Unable to select media route", ex);
             }
@@ -208,9 +219,13 @@ public class MediaRouter2Manager {
      * @param packageName the package name of the application that should stop routing
      */
     public void unselectRoute(@NonNull String packageName) {
-        if (mClient != null) {
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
             try {
-                mMediaRouterService.selectClientRoute2(mClient, packageName, null);
+                mMediaRouterService.selectClientRoute2(client, packageName, null);
             } catch (RemoteException ex) {
                 Log.e(TAG, "Unable to select media route", ex);
             }
@@ -227,10 +242,6 @@ public class MediaRouter2Manager {
         return -1;
     }
 
-    MediaRoute2ProviderInfo getProvider(int index) {
-        return mProviders.get(index);
-    }
-
     void updateProvider(@NonNull MediaRoute2ProviderInfo provider) {
         if (provider == null || !provider.isValid()) {
             Log.w(TAG, "Ignoring invalid provider : " + provider);
@@ -241,7 +252,7 @@ public class MediaRouter2Manager {
 
         final int index = findProviderIndex(provider);
         if (index >= 0) {
-            final MediaRoute2ProviderInfo prevProvider = getProvider(index);
+            final MediaRoute2ProviderInfo prevProvider = mProviders.get(index);
             final Set<String> updatedRouteIds = new HashSet<>();
             for (MediaRoute2Info routeInfo : routes) {
                 final MediaRoute2Info prevRoute = prevProvider.getRoute(routeInfo.getId());
@@ -374,11 +385,9 @@ public class MediaRouter2Manager {
         }
 
         void notifyRoutes() {
-            for (MediaRoute2ProviderInfo provider : mProviders) {
-                for (MediaRoute2Info routeInfo : provider.getRoutes()) {
-                    mExecutor.execute(
-                            () -> mCallback.onRouteAdded(routeInfo));
-                }
+            for (MediaRoute2Info routeInfo : mRoutes) {
+                mExecutor.execute(
+                        () -> mCallback.onRouteAdded(routeInfo));
             }
         }
     }
diff --git a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
index 4282a5bcea01..946fb5e0e0ec 100644
--- a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
+++ b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
@@ -44,8 +44,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
-import java.util.concurrent.SynchronousQueue;
-import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
@@ -92,9 +91,8 @@ public class MediaRouterManagerTest {
         mContext = InstrumentationRegistry.getTargetContext();
         mManager = MediaRouter2Manager.getInstance(mContext);
         mRouter = MediaRouter2.getInstance(mContext);
-        mExecutor = new ThreadPoolExecutor(
-            1, 20, 3, TimeUnit.SECONDS,
-            new SynchronousQueue<Runnable>());
+        //TODO: If we need to support thread pool executors, change this to thread pool executor.
+        mExecutor = Executors.newSingleThreadExecutor();
         mPackageName = mContext.getPackageName();
     }
 
@@ -116,36 +114,34 @@ public class MediaRouterManagerTest {
     public void testRouteAdded() {
         MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
 
-        mManager.addCallback(mExecutor, mockCallback);
+        mManager.registerCallback(mExecutor, mockCallback);
 
         verify(mockCallback, timeout(TIMEOUT_MS)).onRouteAdded(argThat(
                 (MediaRoute2Info info) ->
                         info.getId().equals(ROUTE_ID1) && info.getName().equals(ROUTE_NAME1)));
-        mManager.removeCallback(mockCallback);
+        mManager.unregisterCallback(mockCallback);
     }
 
     //TODO: Recover this test when media router 2 is finalized.
     public void testRouteRemoved() {
         MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
-        mManager.addCallback(mExecutor, mockCallback);
+        mManager.registerCallback(mExecutor, mockCallback);
 
         MediaRouter2.Callback mockRouterCallback = mock(MediaRouter2.Callback.class);
 
         //TODO: Figure out a more proper way to test.
         // (Control requests shouldn't be used in this way.)
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                (Runnable) () -> {
-                    mRouter.addCallback(CONTROL_CATEGORIES_ALL, mExecutor, mockRouterCallback);
-                    mRouter.sendControlRequest(
-                            new MediaRoute2Info.Builder(ROUTE_ID2, ROUTE_NAME2).build(),
-                            new Intent(ACTION_REMOVE_ROUTE));
-                    mRouter.removeCallback(mockRouterCallback);
-                }
-        );
+        mRouter.setControlCategories(CONTROL_CATEGORIES_ALL);
+        mRouter.registerCallback(mExecutor, mockRouterCallback);
+        mRouter.sendControlRequest(
+                new MediaRoute2Info.Builder(ROUTE_ID2, ROUTE_NAME2).build(),
+                new Intent(ACTION_REMOVE_ROUTE));
+        mRouter.unregisterCallback(mockRouterCallback);
+
         verify(mockCallback, timeout(TIMEOUT_MS)).onRouteRemoved(argThat(
                 (MediaRoute2Info info) ->
                         info.getId().equals(ROUTE_ID2) && info.getName().equals(ROUTE_NAME2)));
-        mManager.removeCallback(mockCallback);
+        mManager.unregisterCallback(mockCallback);
     }
 
     /**
@@ -154,17 +150,14 @@ public class MediaRouterManagerTest {
     @Test
     public void testControlCategory() throws Exception {
         MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
-        mManager.addCallback(mExecutor, mockCallback);
+        mManager.registerCallback(mExecutor, mockCallback);
 
         MediaRouter2.Callback mockRouterCallback = mock(MediaRouter2.Callback.class);
 
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> {
-                    mRouter.addCallback(CONTROL_CATEGORIES_SPECIAL,
-                            mExecutor, mockRouterCallback);
-                    mRouter.removeCallback(mockRouterCallback);
-                }
-        );
+        mRouter.setControlCategories(CONTROL_CATEGORIES_SPECIAL);
+        mRouter.registerCallback(mExecutor, mockRouterCallback);
+        mRouter.unregisterCallback(mockRouterCallback);
+
         verify(mockCallback, timeout(TIMEOUT_MS))
                 .onRouteListChanged(argThat(routes -> routes.size() > 0));
 
@@ -174,7 +167,7 @@ public class MediaRouterManagerTest {
         Assert.assertEquals(1, routes.size());
         Assert.assertNotNull(routes.get(ROUTE_ID_SPECIAL_CATEGORY));
 
-        mManager.removeCallback(mockCallback);
+        mManager.unregisterCallback(mockCallback);
     }
 
     @Test
@@ -182,11 +175,7 @@ public class MediaRouterManagerTest {
         CountDownLatch latch = new CountDownLatch(1);
 
         MediaRouter2.Callback mockRouterCallback = mock(MediaRouter2.Callback.class);
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> {
-                    mRouter.addCallback(CONTROL_CATEGORIES_ALL, mExecutor, mockRouterCallback);
-                }
-        );
+        mRouter.registerCallback(mExecutor, mockRouterCallback);
 
         MediaRouter2Manager.Callback managerCallback = new MediaRouter2Manager.Callback() {
             MediaRoute2Info mSelectedRoute = null;
@@ -210,11 +199,11 @@ public class MediaRouterManagerTest {
             }
         };
 
-        mManager.addCallback(mExecutor, managerCallback);
+        mManager.registerCallback(mExecutor, managerCallback);
 
         Assert.assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
 
-        mManager.removeCallback(managerCallback);
+        mManager.unregisterCallback(managerCallback);
     }
 
     /**
@@ -225,12 +214,10 @@ public class MediaRouterManagerTest {
         MediaRouter2Manager.Callback managerCallback = mock(MediaRouter2Manager.Callback.class);
         MediaRouter2.Callback routerCallback = mock(MediaRouter2.Callback.class);
 
-        mManager.addCallback(mExecutor, managerCallback);
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> {
-                    mRouter.addCallback(CONTROL_CATEGORIES_ALL, mExecutor, routerCallback);
-                }
-        );
+        mManager.registerCallback(mExecutor, managerCallback);
+        mRouter.setControlCategories(CONTROL_CATEGORIES_ALL);
+        mRouter.registerCallback(mExecutor, routerCallback);
+
         verify(managerCallback, timeout(TIMEOUT_MS))
                 .onRouteListChanged(argThat(routes -> routes.size() > 0));
 
@@ -253,12 +240,8 @@ public class MediaRouterManagerTest {
                 .onRouteChanged(argThat(routeInfo -> TextUtils.equals(ROUTE_ID2, routeInfo.getId())
                         && TextUtils.equals(routeInfo.getClientPackageName(), null)));
 
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> {
-                    mRouter.removeCallback(routerCallback);
-                }
-        );
-        mManager.removeCallback(managerCallback);
+        mRouter.unregisterCallback(routerCallback);
+        mManager.unregisterCallback(managerCallback);
     }
 
     Map<String, MediaRoute2Info> createRouteMap(List<MediaRoute2Info> routes) {