When posture changes, re-register prox listeners

Listeners that were skipped from registering
because the previous posture didn't support prox
need to be re-registered.

Test: atest PostureDependentProximitySensorTest
Test: enter AOD from the unfolded state and then
close the phone. From AOD in the folded state,
cover the prox sensor => observe that AOD display
turns off after 10+ seconds.
Fixes: 272691904

Change-Id: Ib8a72c2af01a546019ef4a611ca9877e20881346
diff --git a/packages/SystemUI/src/com/android/systemui/util/sensors/PostureDependentProximitySensor.java b/packages/SystemUI/src/com/android/systemui/util/sensors/PostureDependentProximitySensor.java
index 460b7d9..a5828c7 100644
--- a/packages/SystemUI/src/com/android/systemui/util/sensors/PostureDependentProximitySensor.java
+++ b/packages/SystemUI/src/com/android/systemui/util/sensors/PostureDependentProximitySensor.java
@@ -23,6 +23,8 @@
 import com.android.systemui.util.concurrency.DelayableExecutor;
 import com.android.systemui.util.concurrency.Execution;
 
+import java.util.HashSet;
+
 import javax.inject.Inject;
 
 /**
@@ -37,6 +39,7 @@
     private final ThresholdSensor[] mPostureToPrimaryProxSensorMap;
     private final ThresholdSensor[] mPostureToSecondaryProxSensorMap;
 
+    private final HashSet<Listener> mListenersRegisteredWhenProxUnavailable = new HashSet<>();
     private final DevicePostureController mDevicePostureController;
 
     @Inject
@@ -69,6 +72,25 @@
         mDevicePostureController.removeCallback(mDevicePostureCallback);
     }
 
+    @Override
+    public void register(ThresholdSensor.Listener listener) {
+        if (!isLoaded()) {
+            logDebug("No prox sensor when registering listener=" + listener);
+            mListenersRegisteredWhenProxUnavailable.add(listener);
+        }
+
+        super.register(listener);
+    }
+
+    @Override
+    public void unregister(ThresholdSensor.Listener listener) {
+        if (mListenersRegisteredWhenProxUnavailable.remove(listener)) {
+            logDebug("Removing listener from mListenersRegisteredWhenProxUnavailable "
+                    + listener);
+        }
+        super.unregister(listener);
+    }
+
     private void chooseSensors() {
         if (mDevicePosture >= mPostureToPrimaryProxSensorMap.length
                 || mDevicePosture >= mPostureToSecondaryProxSensorMap.length) {
@@ -98,6 +120,14 @@
 
             mInitializedListeners = false;
             registerInternal();
+
+            final Listener[] listenersToReregister =
+                    mListenersRegisteredWhenProxUnavailable.toArray(new Listener[0]);
+            mListenersRegisteredWhenProxUnavailable.clear();
+            for (Listener listener : listenersToReregister) {
+                logDebug("Re-register listener " + listener);
+                register(listener);
+            }
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/sensors/PostureDependentProximitySensorTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/sensors/PostureDependentProximitySensorTest.java
index 075f393..84129be 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/sensors/PostureDependentProximitySensorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/sensors/PostureDependentProximitySensorTest.java
@@ -16,10 +16,16 @@
 
 package com.android.systemui.util.sensors;
 
+import static com.android.systemui.statusbar.policy.DevicePostureController.DEVICE_POSTURE_CLOSED;
+
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
 import android.content.res.Resources;
+import android.hardware.Sensor;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
@@ -46,28 +52,52 @@
     @Mock private Resources mResources;
     @Mock private DevicePostureController mDevicePostureController;
     @Mock private AsyncSensorManager mSensorManager;
+    @Mock private Sensor mMockedPrimaryProxSensor;
 
     @Captor private ArgumentCaptor<DevicePostureController.Callback> mPostureListenerCaptor =
             ArgumentCaptor.forClass(DevicePostureController.Callback.class);
     private DevicePostureController.Callback mPostureListener;
 
-    private PostureDependentProximitySensor mProximitySensor;
-    private FakeExecutor mFakeExecutor = new FakeExecutor(new FakeSystemClock());
+    private PostureDependentProximitySensor mPostureDependentProximitySensor;
+    private ThresholdSensor[] mPrimaryProxSensors;
+    private ThresholdSensor[] mSecondaryProxSensors;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         allowTestableLooperAsMainThread();
 
-        mProximitySensor = new PostureDependentProximitySensor(
-                new ThresholdSensor[DevicePostureController.SUPPORTED_POSTURES_SIZE],
-                new ThresholdSensor[DevicePostureController.SUPPORTED_POSTURES_SIZE],
-                mFakeExecutor,
+        setupProximitySensors(DEVICE_POSTURE_CLOSED);
+        mPostureDependentProximitySensor = new PostureDependentProximitySensor(
+                mPrimaryProxSensors,
+                mSecondaryProxSensors,
+                new FakeExecutor(new FakeSystemClock()),
                 new FakeExecution(),
                 mDevicePostureController
         );
     }
 
+    /**
+     * Support a proximity sensor only for the given devicePosture for the primary sensor.
+     * Otherwise, all other postures don't support prox.
+     */
+    private void setupProximitySensors(
+            @DevicePostureController.DevicePostureInt int proxExistsForPosture) {
+        final ThresholdSensorImpl.Builder sensorBuilder = new ThresholdSensorImpl.BuilderFactory(
+                mResources, mSensorManager, new FakeExecution()).createBuilder();
+
+        mPrimaryProxSensors = new ThresholdSensor[DevicePostureController.SUPPORTED_POSTURES_SIZE];
+        mSecondaryProxSensors =
+                new ThresholdSensor[DevicePostureController.SUPPORTED_POSTURES_SIZE];
+        for (int i = 0; i < DevicePostureController.SUPPORTED_POSTURES_SIZE; i++) {
+            mPrimaryProxSensors[i] = sensorBuilder.setSensor(null).setThresholdValue(0).build();
+            mSecondaryProxSensors[i] = sensorBuilder.setSensor(null).setThresholdValue(0).build();
+        }
+
+        mPrimaryProxSensors[proxExistsForPosture] = sensorBuilder
+                .setSensor(mMockedPrimaryProxSensor).setThresholdValue(5).build();
+    }
+
     @Test
     public void testPostureChangeListenerAdded() {
         capturePostureListener();
@@ -83,30 +113,59 @@
 
         // THEN device posture is updated to DEVICE_POSTURE_OPENED
         assertEquals(DevicePostureController.DEVICE_POSTURE_OPENED,
-                mProximitySensor.mDevicePosture);
+                mPostureDependentProximitySensor.mDevicePosture);
 
         // WHEN the posture changes to DEVICE_POSTURE_CLOSED
-        mPostureListener.onPostureChanged(DevicePostureController.DEVICE_POSTURE_CLOSED);
+        mPostureListener.onPostureChanged(DEVICE_POSTURE_CLOSED);
 
         // THEN device posture is updated to DEVICE_POSTURE_CLOSED
-        assertEquals(DevicePostureController.DEVICE_POSTURE_CLOSED,
-                mProximitySensor.mDevicePosture);
+        assertEquals(DEVICE_POSTURE_CLOSED,
+                mPostureDependentProximitySensor.mDevicePosture);
 
         // WHEN the posture changes to DEVICE_POSTURE_FLIPPED
         mPostureListener.onPostureChanged(DevicePostureController.DEVICE_POSTURE_FLIPPED);
 
         // THEN device posture is updated to DEVICE_POSTURE_FLIPPED
         assertEquals(DevicePostureController.DEVICE_POSTURE_FLIPPED,
-                mProximitySensor.mDevicePosture);
+                mPostureDependentProximitySensor.mDevicePosture);
 
         // WHEN the posture changes to DEVICE_POSTURE_HALF_OPENED
         mPostureListener.onPostureChanged(DevicePostureController.DEVICE_POSTURE_HALF_OPENED);
 
         // THEN device posture is updated to DEVICE_POSTURE_HALF_OPENED
         assertEquals(DevicePostureController.DEVICE_POSTURE_HALF_OPENED,
-                mProximitySensor.mDevicePosture);
+                mPostureDependentProximitySensor.mDevicePosture);
     }
 
+    @Test
+    public void proxSensorRegisters_proxSensorValid() {
+        // GIVEN posture that supports a valid posture with a prox sensor
+        capturePostureListener();
+        mPostureListener.onPostureChanged(DEVICE_POSTURE_CLOSED);
+
+        // WHEN a listener registers
+        mPostureDependentProximitySensor.register(mock(ThresholdSensor.Listener.class));
+
+        // THEN PostureDependentProximitySensor is registered
+        assertTrue(mPostureDependentProximitySensor.isRegistered());
+    }
+
+    @Test
+    public void proxSensorReregisters_postureChangesAndNewlySupportsProx() {
+        // GIVEN there's a registered listener but posture doesn't support prox
+        assertFalse(mPostureDependentProximitySensor.isRegistered());
+        mPostureDependentProximitySensor.register(mock(ThresholdSensor.Listener.class));
+        assertFalse(mPostureDependentProximitySensor.isRegistered());
+
+        // WHEN posture that supports a valid posture with a prox sensor
+        capturePostureListener();
+        mPostureListener.onPostureChanged(DEVICE_POSTURE_CLOSED);
+
+        // THEN PostureDependentProximitySensor is registered
+        assertTrue(mPostureDependentProximitySensor.isRegistered());
+    }
+
+
     private void capturePostureListener() {
         verify(mDevicePostureController).addCallback(mPostureListenerCaptor.capture());
         mPostureListener = mPostureListenerCaptor.getValue();