New system for plugin + tuner integrations called extensions

An ExtensionController provides an easy way to say I need an
object of interface X. Then a plugin or a tuner factory can
actually provide X when needed or fallback to a default implementation.

Test: runtest systemui
Change-Id: I5e1b76def3c790d7f673867648ffeb13c4d0a829

9 files changed, 579 insertions, 93 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/Dependency.java b/packages/SystemUI/src/com/android/systemui/Dependency.java
index f1e7d53b037f..ac7ab9de7f11 100644
--- a/packages/SystemUI/src/com/android/systemui/Dependency.java
+++ b/packages/SystemUI/src/com/android/systemui/Dependency.java
@@ -47,6 +47,8 @@ import com.android.systemui.statusbar.policy.DarkIconDispatcher;
 import com.android.systemui.statusbar.policy.DataSaverController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedControllerImpl;
+import com.android.systemui.statusbar.policy.ExtensionController;
+import com.android.systemui.statusbar.policy.ExtensionControllerImpl;
 import com.android.systemui.statusbar.policy.FlashlightController;
 import com.android.systemui.statusbar.policy.FlashlightControllerImpl;
 import com.android.systemui.statusbar.policy.HotspotController;
@@ -233,6 +235,9 @@ public class Dependency extends SystemUI {
         mProviders.put(FragmentService.class, () ->
                 new FragmentService(mContext));
 
+        mProviders.put(ExtensionController.class, () ->
+                new ExtensionControllerImpl());
+
         // Put all dependencies above here so the factory can override them if it wants.
         SystemUIFactory.getInstance().injectDependencies(mProviders, mContext);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/plugins/PluginManager.java b/packages/SystemUI/src/com/android/systemui/plugins/PluginManager.java
index 8b4bd7bc9f7b..0c3e40c7d891 100644
--- a/packages/SystemUI/src/com/android/systemui/plugins/PluginManager.java
+++ b/packages/SystemUI/src/com/android/systemui/plugins/PluginManager.java
@@ -133,14 +133,7 @@ public class PluginManager extends BroadcastReceiver {
 
     public <T extends Plugin> void addPluginListener(PluginListener<T> listener, Class<?> cls,
             boolean allowMultiple) {
-        ProvidesInterface info = cls.getDeclaredAnnotation(ProvidesInterface.class);
-        if (info == null) {
-            throw new RuntimeException(cls + " doesn't provide an interface");
-        }
-        if (TextUtils.isEmpty(info.action())) {
-            throw new RuntimeException(cls + " doesn't provide an action");
-        }
-        addPluginListener(info.action(), listener, cls, allowMultiple);
+        addPluginListener(getAction(cls), listener, cls, allowMultiple);
     }
 
     public <T extends Plugin> void addPluginListener(String action, PluginListener<T> listener,
@@ -285,6 +278,17 @@ public class PluginManager extends BroadcastReceiver {
         return new PluginContextWrapper(mContext.createApplicationContext(info, 0), classLoader);
     }
 
+    public static <P> String getAction(Class<P> cls) {
+        ProvidesInterface info = cls.getDeclaredAnnotation(ProvidesInterface.class);
+        if (info == null) {
+            throw new RuntimeException(cls + " doesn't provide an interface");
+        }
+        if (TextUtils.isEmpty(info.action())) {
+            throw new RuntimeException(cls + " doesn't provide an action");
+        }
+        return info.action();
+    }
+
     private class AllPluginClassLoader extends ClassLoader {
         public AllPluginClassLoader(ClassLoader classLoader) {
             super(classLoader);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardBottomAreaView.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardBottomAreaView.java
index 8f63d45e04cd..2538bdd79b0e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardBottomAreaView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardBottomAreaView.java
@@ -19,9 +19,10 @@ package com.android.systemui.statusbar.phone;
 import static android.view.accessibility.AccessibilityNodeInfo.ACTION_CLICK;
 import static android.view.accessibility.AccessibilityNodeInfo.AccessibilityAction;
 
+import static com.android.systemui.tuner.LockscreenFragment.LOCKSCREEN_LEFT_BUTTON;
 import static com.android.systemui.tuner.LockscreenFragment.LOCKSCREEN_LEFT_UNLOCK;
+import static com.android.systemui.tuner.LockscreenFragment.LOCKSCREEN_RIGHT_BUTTON;
 import static com.android.systemui.tuner.LockscreenFragment.LOCKSCREEN_RIGHT_UNLOCK;
-import static com.android.systemui.tuner.LockscreenFragment.getIntentButton;
 
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
@@ -79,9 +80,12 @@ import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.KeyguardAffordanceView;
 import com.android.systemui.statusbar.KeyguardIndicationController;
 import com.android.systemui.statusbar.policy.AccessibilityController;
+import com.android.systemui.statusbar.policy.ExtensionController;
+import com.android.systemui.statusbar.policy.ExtensionController.Extension;
 import com.android.systemui.statusbar.policy.FlashlightController;
 import com.android.systemui.statusbar.policy.PreviewInflater;
 import com.android.systemui.tuner.LockscreenFragment;
+import com.android.systemui.tuner.LockscreenFragment.LockButtonFactory;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.tuner.TunerService.Tunable;
 
@@ -91,8 +95,7 @@ import com.android.systemui.tuner.TunerService.Tunable;
  */
 public class KeyguardBottomAreaView extends FrameLayout implements View.OnClickListener,
         UnlockMethodCache.OnUnlockMethodChangedListener,
-        AccessibilityController.AccessibilityStateChangedCallback, View.OnLongClickListener,
-        Tunable {
+        AccessibilityController.AccessibilityStateChangedCallback, View.OnLongClickListener {
 
     final static String TAG = "StatusBar/KeyguardBottomAreaView";
 
@@ -159,12 +162,10 @@ public class KeyguardBottomAreaView extends FrameLayout implements View.OnClickL
     private Drawable mLeftAssistIcon;
 
     private IntentButton mRightButton = new DefaultRightButton();
-    private IntentButton mRightDefault = mRightButton;
-    private IntentButton mRightPlugin;
+    private Extension<IntentButton> mRightExtension;
     private String mRightButtonStr;
     private IntentButton mLeftButton = new DefaultLeftButton();
-    private IntentButton mLeftDefault = mLeftButton;
-    private IntentButton mLeftPlugin;
+    private Extension<IntentButton> mLeftExtension;
     private String mLeftButtonStr;
     private LockscreenGestureLogger mLockscreenGestureLogger = new LockscreenGestureLogger();
     private boolean mDozing;
@@ -261,21 +262,28 @@ public class KeyguardBottomAreaView extends FrameLayout implements View.OnClickL
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
         mAccessibilityController.addStateChangedCallback(this);
-        Dependency.get(PluginManager.class).addPluginListener(RIGHT_BUTTON_PLUGIN,
-                mRightListener, IntentButtonProvider.class, false /* Only allow one */);
-        Dependency.get(PluginManager.class).addPluginListener(LEFT_BUTTON_PLUGIN,
-                mLeftListener, IntentButtonProvider.class, false /* Only allow one */);
-        Dependency.get(TunerService.class).addTunable(this, LockscreenFragment.LOCKSCREEN_LEFT_BUTTON,
-                LockscreenFragment.LOCKSCREEN_RIGHT_BUTTON);
+        mRightExtension = Dependency.get(ExtensionController.class).newExtension(IntentButton.class)
+                .withPlugin(IntentButtonProvider.class, RIGHT_BUTTON_PLUGIN,
+                        p -> p.getIntentButton())
+                .withTunerFactory(new LockButtonFactory(mContext, LOCKSCREEN_RIGHT_BUTTON))
+                .withDefault(() -> new DefaultRightButton())
+                .withCallback(button -> setRightButton(button))
+                .build();
+        mLeftExtension = Dependency.get(ExtensionController.class).newExtension(IntentButton.class)
+                .withPlugin(IntentButtonProvider.class, LEFT_BUTTON_PLUGIN,
+                        p -> p.getIntentButton())
+                .withTunerFactory(new LockButtonFactory(mContext, LOCKSCREEN_LEFT_BUTTON))
+                .withDefault(() -> new DefaultLeftButton())
+                .withCallback(button -> setLeftButton(button))
+                .build();
     }
 
     @Override
     protected void onDetachedFromWindow() {
         super.onDetachedFromWindow();
         mAccessibilityController.removeStateChangedCallback(this);
-        Dependency.get(PluginManager.class).removePluginListener(mRightListener);
-        Dependency.get(PluginManager.class).removePluginListener(mLeftListener);
-        Dependency.get(TunerService.class).removeTunable(this);
+        mRightExtension.destroy();
+        mLeftExtension.destroy();
     }
 
     private void initAccessibility() {
@@ -790,63 +798,21 @@ public class KeyguardBottomAreaView extends FrameLayout implements View.OnClickL
         inflateCameraPreview();
     }
 
-    @Override
-    public void onTuningChanged(String key, String newValue) {
-        if (LockscreenFragment.LOCKSCREEN_LEFT_BUTTON.equals(key)) {
-            mLeftButtonStr = newValue;
-            mLeftIsVoiceAssist = TextUtils.isEmpty(mLeftButtonStr) && mLeftPlugin == null;
-            mLeftButton = getIntentButton(mContext, mLeftButtonStr, mLeftPlugin, mLeftDefault);
-            updateLeftAffordance();
-        } else {
-            mRightButtonStr = newValue;
-            mRightButton = getIntentButton(mContext, mRightButtonStr, mRightPlugin, mRightDefault);
-            updateRightAffordanceIcon();
-            updateCameraVisibility();
-            inflateCameraPreview();
-        }
-    }
-
     private void setRightButton(IntentButton button) {
-        mRightPlugin = button;
-        mRightButton = getIntentButton(mContext, mRightButtonStr, mRightPlugin, mRightDefault);
+        mRightButton = button;
         updateRightAffordanceIcon();
         updateCameraVisibility();
         inflateCameraPreview();
     }
 
     private void setLeftButton(IntentButton button) {
-        mLeftPlugin = button;
-        mLeftButton = getIntentButton(mContext, mLeftButtonStr, mLeftPlugin, mLeftDefault);
-        mLeftIsVoiceAssist = false;
+        mLeftButton = button;
+        if (!(mLeftButton instanceof DefaultLeftButton)) {
+            mLeftIsVoiceAssist = false;
+        }
         updateLeftAffordance();
     }
 
-    private final PluginListener<IntentButtonProvider> mRightListener =
-            new PluginListener<IntentButtonProvider>() {
-        @Override
-        public void onPluginConnected(IntentButtonProvider plugin, Context pluginContext) {
-            setRightButton(plugin.getIntentButton());
-        }
-
-        @Override
-        public void onPluginDisconnected(IntentButtonProvider plugin) {
-            setRightButton(null);
-        }
-    };
-
-    private final PluginListener<IntentButtonProvider> mLeftListener =
-            new PluginListener<IntentButtonProvider>() {
-        @Override
-        public void onPluginConnected(IntentButtonProvider plugin, Context pluginContext) {
-            setLeftButton(plugin.getIntentButton());
-        }
-
-        @Override
-        public void onPluginDisconnected(IntentButtonProvider plugin) {
-            setLeftButton(null);
-        }
-    };
-
     public void setDozing(boolean dozing, boolean animate) {
         mDozing = dozing;
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionController.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionController.java
new file mode 100644
index 000000000000..eaf89259ed59
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionController.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+import com.android.systemui.Dependency;
+import com.android.systemui.plugins.Plugin;
+
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+/**
+ * Utility class used to select between a plugin, tuner settings, and a default implementation
+ * of an interface.
+ */
+public interface ExtensionController {
+
+    <T> ExtensionBuilder<T> newExtension(Class<T> cls);
+
+    interface Extension<T> {
+        T get();
+        void destroy();
+    }
+
+    interface ExtensionBuilder<T> {
+        ExtensionBuilder<T> withTunerFactory(TunerFactory<T> factory);
+        <P extends T> ExtensionBuilder<T> withPlugin(Class<P> cls);
+        <P extends T> ExtensionBuilder<T> withPlugin(Class<P> cls, String action);
+        <P> ExtensionBuilder<T> withPlugin(Class<P> cls, String action,
+                PluginConverter<T, P> converter);
+        ExtensionBuilder<T> withDefault(Supplier<T> def);
+        ExtensionBuilder<T> withCallback(Consumer<T> callback);
+        Extension build();
+    }
+
+    public interface PluginConverter<T, P> {
+        T getInterfaceFromPlugin(P plugin);
+    }
+
+    public interface TunerFactory<T> {
+        String[] keys();
+        T create(Map<String, String> settings);
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionControllerImpl.java
new file mode 100644
index 000000000000..fefbaa31af2a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ExtensionControllerImpl.java
@@ -0,0 +1,236 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+import com.android.systemui.Dependency;
+import com.android.systemui.plugins.Plugin;
+import com.android.systemui.plugins.PluginListener;
+import com.android.systemui.plugins.PluginManager;
+import com.android.systemui.tuner.TunerService;
+import com.android.systemui.tuner.TunerService.Tunable;
+
+import android.content.Context;
+import android.util.ArrayMap;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+public class ExtensionControllerImpl implements ExtensionController {
+
+    @Override
+    public <T> ExtensionBuilder<T> newExtension(Class<T> cls) {
+        return new ExtensionBuilder<>();
+    }
+
+    private interface Producer<T> {
+        T get();
+        void destroy();
+    }
+
+    private class ExtensionBuilder<T> implements ExtensionController.ExtensionBuilder<T> {
+
+        private ExtensionImpl<T> mExtension = new ExtensionImpl<>();
+
+        @Override
+        public ExtensionController.ExtensionBuilder<T> withTunerFactory(TunerFactory<T> factory) {
+            mExtension.addTunerFactory(factory, factory.keys());
+            return this;
+        }
+
+        @Override
+        public <P extends T> ExtensionController.ExtensionBuilder<T> withPlugin(Class<P> cls) {
+            return withPlugin(cls, PluginManager.getAction(cls));
+        }
+
+        @Override
+        public <P extends T> ExtensionController.ExtensionBuilder<T> withPlugin(Class<P> cls,
+                String action) {
+            return withPlugin(cls, action, null);
+        }
+
+        @Override
+        public <P> ExtensionController.ExtensionBuilder<T> withPlugin(Class<P> cls,
+                String action, PluginConverter<T, P> converter) {
+            mExtension.addPlugin(action, cls, converter);
+            return this;
+        }
+
+        @Override
+        public ExtensionController.ExtensionBuilder<T> withDefault(Supplier<T> def) {
+            mExtension.addDefault(def);
+            return this;
+        }
+
+        @Override
+        public ExtensionController.ExtensionBuilder<T> withCallback(
+                Consumer<T> callback) {
+            mExtension.mCallbacks.add(callback);
+            return this;
+        }
+
+        @Override
+        public ExtensionController.Extension build() {
+            // Manually sort, plugins first, tuners second, defaults last.
+            Collections.sort(mExtension.mProducers, (o1, o2) -> {
+                if (o1 instanceof ExtensionImpl.PluginItem) {
+                    if (o2 instanceof ExtensionImpl.PluginItem) {
+                        return 0;
+                    } else {
+                        return -1;
+                    }
+                }
+                if (o1 instanceof ExtensionImpl.TunerItem) {
+                    if (o2 instanceof ExtensionImpl.PluginItem) {
+                        return 1;
+                    } else if (o2 instanceof ExtensionImpl.TunerItem) {
+                        return 0;
+                    } else {
+                        return -1;
+                    }
+                }
+                return 0;
+            });
+            mExtension.notifyChanged();
+            return mExtension;
+        }
+    }
+
+    private class ExtensionImpl<T> implements ExtensionController.Extension<T> {
+        private final ArrayList<Producer<T>> mProducers = new ArrayList<>();
+        private final ArrayList<Consumer<T>> mCallbacks = new ArrayList<>();
+        private T mItem;
+
+        @Override
+        public T get() {
+            return mItem;
+        }
+
+        @Override
+        public void destroy() {
+            for (int i = 0; i < mProducers.size(); i++) {
+                mProducers.get(i).destroy();
+            }
+        }
+
+        private void notifyChanged() {
+            for (int i = 0; i < mProducers.size(); i++) {
+                final T item = mProducers.get(i).get();
+                if (item != null) {
+                    mItem = item;
+                    break;
+                }
+            }
+            for (int i = 0; i < mCallbacks.size(); i++) {
+                mCallbacks.get(i).accept(mItem);
+            }
+        }
+
+        public void addDefault(Supplier<T> def) {
+            mProducers.add(new Default(def));
+        }
+
+        public <P> void addPlugin(String action, Class<P> cls, PluginConverter<T, P> converter) {
+            mProducers.add(new PluginItem(action, cls, converter));
+        }
+
+        public void addTunerFactory(TunerFactory<T> factory, String[] keys) {
+            mProducers.add(new TunerItem(factory, factory.keys()));
+        }
+
+        private class PluginItem<P extends Plugin> implements Producer<T>, PluginListener<P> {
+            private final PluginConverter<T, P> mConverter;
+            private T mItem;
+
+            public PluginItem(String action, Class<P> cls, PluginConverter<T, P> converter) {
+                mConverter = converter;
+                Dependency.get(PluginManager.class).addPluginListener(action, this, cls);
+            }
+
+            @Override
+            public void onPluginConnected(P plugin, Context pluginContext) {
+                if (mConverter != null) {
+                    mItem = mConverter.getInterfaceFromPlugin(plugin);
+                } else {
+                    mItem = (T) plugin;
+                }
+                notifyChanged();
+            }
+
+            @Override
+            public void onPluginDisconnected(P plugin) {
+                mItem = null;
+                notifyChanged();
+            }
+
+            @Override
+            public T get() {
+                return mItem;
+            }
+
+            @Override
+            public void destroy() {
+                Dependency.get(PluginManager.class).removePluginListener(this);
+            }
+        }
+
+        private class TunerItem<T> implements Producer<T>, Tunable {
+            private final TunerFactory<T> mFactory;
+            private final ArrayMap<String, String> mSettings = new ArrayMap<>();
+            private T mItem;
+
+            public TunerItem(TunerFactory<T> factory, String... setting) {
+                mFactory = factory;
+                Dependency.get(TunerService.class).addTunable(this, setting);
+            }
+
+            @Override
+            public T get() {
+                return mItem;
+            }
+
+            @Override
+            public void destroy() {
+                Dependency.get(TunerService.class).removeTunable(this);
+            }
+
+            @Override
+            public void onTuningChanged(String key, String newValue) {
+                mSettings.put(key, newValue);
+                mItem = mFactory.create(mSettings);
+                notifyChanged();
+            }
+        }
+
+        private class Default<T> implements Producer<T> {
+            private final Supplier<T> mSupplier;
+
+            public Default(Supplier<T> supplier) {
+                mSupplier = supplier;
+            }
+
+            @Override
+            public T get() {
+                return mSupplier.get();
+            }
+
+            @Override
+            public void destroy() {
+
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/tuner/LockscreenFragment.java b/packages/SystemUI/src/com/android/systemui/tuner/LockscreenFragment.java
index 410d3d2e1b53..2df1793acb47 100644
--- a/packages/SystemUI/src/com/android/systemui/tuner/LockscreenFragment.java
+++ b/packages/SystemUI/src/com/android/systemui/tuner/LockscreenFragment.java
@@ -52,11 +52,13 @@ import com.android.systemui.R;
 import com.android.systemui.plugins.IntentButtonProvider.IntentButton;
 import com.android.systemui.statusbar.ScalingDrawableWrapper;
 import com.android.systemui.statusbar.phone.ExpandableIndicator;
+import com.android.systemui.statusbar.policy.ExtensionController.TunerFactory;
 import com.android.systemui.tuner.ShortcutParser.Shortcut;
 import com.android.systemui.tuner.TunerService.Tunable;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Consumer;
 
 public class LockscreenFragment extends PreferenceFragment {
@@ -313,26 +315,39 @@ public class LockscreenFragment extends PreferenceFragment {
         }
     }
 
-    public static IntentButton getIntentButton(Context context, String buttonStr,
-            IntentButton plugin, IntentButton def) {
-        // Plugin wins.
-        if (plugin != null) return plugin;
-        // Then tuner options.
-        if (!TextUtils.isEmpty(buttonStr)) {
-            if (buttonStr.contains("::")) {
-                Shortcut shortcut = getShortcutInfo(context, buttonStr);
-                if (shortcut != null) {
-                    return new ShortcutButton(context, shortcut);
-                }
-            } else if (buttonStr.contains("/")) {
-                ActivityInfo info = getActivityinfo(context, buttonStr);
-                if (info != null) {
-                    return new ActivityButton(context, info);
+    public static class LockButtonFactory implements TunerFactory<IntentButton> {
+
+        private final String mKey;
+        private final Context mContext;
+
+        public LockButtonFactory(Context context, String key) {
+            mContext = context;
+            mKey = key;
+        }
+
+        @Override
+        public String[] keys() {
+            return new String[]{mKey};
+        }
+
+        @Override
+        public IntentButton create(Map<String, String> settings) {
+            String buttonStr = settings.get(mKey);
+            if (!TextUtils.isEmpty(buttonStr)) {
+                if (buttonStr.contains("::")) {
+                    Shortcut shortcut = getShortcutInfo(mContext, buttonStr);
+                    if (shortcut != null) {
+                        return new ShortcutButton(mContext, shortcut);
+                    }
+                } else if (buttonStr.contains("/")) {
+                    ActivityInfo info = getActivityinfo(mContext, buttonStr);
+                    if (info != null) {
+                        return new ActivityButton(mContext, info);
+                    }
                 }
             }
+            return null;
         }
-        // Then default.
-        return def;
     }
 
     private static class ShortcutButton implements IntentButton {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
index 6c454e126eaf..c0e7e8061407 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
@@ -23,7 +23,6 @@ import android.os.Looper;
 import android.os.MessageQueue;
 import android.support.test.InstrumentationRegistry;
 import android.util.ArrayMap;
-import android.util.Log;
 
 import com.android.systemui.Dependency.DependencyKey;
 import com.android.systemui.utils.TestableContext;
@@ -32,8 +31,6 @@ import com.android.systemui.utils.leaks.Tracker;
 import org.junit.After;
 import org.junit.Before;
 
-import java.lang.Thread.UncaughtExceptionHandler;
-
 /**
  * Base class that does System UI specific setup.
  */
@@ -92,8 +89,10 @@ public abstract class SysuiTestCase {
         return null;
     }
 
-    public <T> void injectMockDependency(Class<T> cls) {
-        injectTestDependency(cls, mock(cls));
+    public <T> T injectMockDependency(Class<T> cls) {
+        final T mock = mock(cls);
+        mDependency.injectTestDependency(cls, mock);
+        return mock;
     }
 
     public <T> void injectTestDependency(Class<T> cls, T obj) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ExtensionControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ExtensionControllerTest.java
new file mode 100644
index 000000000000..e3a5ef04bcc8
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ExtensionControllerTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import com.android.systemui.Dependency;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.plugins.OverlayPlugin;
+import com.android.systemui.plugins.PluginManager;
+import com.android.systemui.statusbar.policy.ExtensionController.Extension;
+import com.android.systemui.statusbar.policy.ExtensionController.TunerFactory;
+import com.android.systemui.tuner.TunerService;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Map;
+import java.util.function.Consumer;
+
+public class ExtensionControllerTest extends SysuiTestCase {
+
+    private PluginManager mPluginManager;
+    private TunerService mTunerService;
+    private ExtensionController mExtensionController;
+
+    @Before
+    public void setup() {
+        mPluginManager = injectMockDependency(PluginManager.class);
+        mTunerService = injectMockDependency(TunerService.class);
+        mExtensionController = Dependency.get(ExtensionController.class);
+    }
+
+    @Test
+    public void testPlugin() {
+        Extension ext = mExtensionController.newExtension(OverlayPlugin.class)
+                .withPlugin(OverlayPlugin.class)
+                .build();
+        verify(mPluginManager).addPluginListener(eq(OverlayPlugin.ACTION), any(),
+                eq(OverlayPlugin.class));
+
+        ext.destroy();
+        verify(mPluginManager).removePluginListener(any());
+    }
+
+    @Test
+    public void testTuner() {
+        String[] keys = new String[] { "key1", "key2" };
+        TunerFactory<Object> factory = new ExtensionController.TunerFactory() {
+            @Override
+            public String[] keys() {
+                return keys;
+            }
+
+            @Override
+            public Object create(Map settings) {
+                return null;
+            }
+        };
+        Extension ext = mExtensionController.newExtension(Object.class)
+                .withTunerFactory(factory)
+                .build();
+        verify(mTunerService).addTunable(any(), eq(keys[0]), eq(keys[1]));
+
+        ext.destroy();
+        verify(mTunerService).removeTunable(any());
+    }
+
+    @Test
+    public void testDefault() {
+        Object o = new Object();
+        Extension ext = mExtensionController.newExtension(Object.class)
+                .withDefault(() -> o)
+                .build();
+        assertEquals(o, ext.get());
+    }
+
+    @Test
+    public void testCallback() {
+        Consumer<Object> callback = mock(Consumer.class);
+        final Object o = new Object();
+        mExtensionController.newExtension(Object.class)
+                .withDefault(() -> o)
+                .withCallback(callback)
+                .build();
+        verify(callback).accept(eq(o));
+    }
+
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeExtensionController.java b/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeExtensionController.java
new file mode 100644
index 000000000000..c0f578373fe4
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeExtensionController.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.android.systemui.utils.leaks;
+
+import static org.mockito.Mockito.mock;
+
+import com.android.systemui.statusbar.policy.ExtensionController;
+
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+public class FakeExtensionController implements ExtensionController {
+
+    private final Tracker mTracker;
+
+    public FakeExtensionController(LeakCheckedTest test) {
+        mTracker = test.getTracker("extension");
+    }
+
+    @Override
+    public <T> ExtensionBuilder<T> newExtension(Class<T> cls) {
+        final Object o = new Object();
+        mTracker.getLeakInfo(o).addAllocation(new Throwable());
+        return new FakeExtensionBuilder(o);
+    }
+
+    private class FakeExtensionBuilder<T> implements ExtensionBuilder<T> {
+        private final Object mAllocation;
+
+        public FakeExtensionBuilder(Object o) {
+            mAllocation = o;
+        }
+
+        @Override
+        public ExtensionBuilder<T> withTunerFactory(TunerFactory<T> factory) {
+            return this;
+        }
+
+        @Override
+        public <P extends T> ExtensionBuilder<T> withPlugin(Class<P> cls) {
+            return this;
+        }
+
+        @Override
+        public <P extends T> ExtensionBuilder<T> withPlugin(Class<P> cls, String action) {
+            return this;
+        }
+
+        @Override
+        public <P> ExtensionBuilder<T> withPlugin(Class<P> cls, String action, PluginConverter<T, P> converter) {
+            return this;
+        }
+
+        @Override
+        public ExtensionBuilder<T> withDefault(Supplier<T> def) {
+            return this;
+        }
+
+        @Override
+        public ExtensionBuilder<T> withCallback(Consumer<T> callback) {
+            return this;
+        }
+
+        @Override
+        public Extension build() {
+            return new FakeExtension(mAllocation);
+        }
+    }
+
+    private class FakeExtension<T> implements Extension<T> {
+        private final Object mAllocation;
+
+        public FakeExtension(Object allocation) {
+            mAllocation = allocation;
+        }
+
+        @Override
+        public T get() {
+            // TODO: Support defaults or things.
+            return null;
+        }
+
+        @Override
+        public void destroy() {
+            mTracker.getLeakInfo(mAllocation).clearAllocations();
+        }
+    }
+}