Restore connected state in Wifi picker.

This part of the change reads the connected state from the extras bundle and updates the HotspotNetworkEntry. The entry uses this state to reset the mCalledConnect boolean and trigger the callback.

The connected enum value is defined locally in WifiTrackerLib, but will be moved to the SharedConnectivity API on the next API update.

BUG: 312561529
Change-Id: I6675024c67f926d28544013b23cef88b2d63ce6e

4 files changed, 115 insertions, 9 deletions

diff --git a/libs/WifiTrackerLib/src/com/android/wifitrackerlib/HotspotNetworkEntry.java b/libs/WifiTrackerLib/src/com/android/wifitrackerlib/HotspotNetworkEntry.java
index 7eb225b27..4a1341454 100644
--- a/libs/WifiTrackerLib/src/com/android/wifitrackerlib/HotspotNetworkEntry.java
+++ b/libs/WifiTrackerLib/src/com/android/wifitrackerlib/HotspotNetworkEntry.java
@@ -98,6 +98,8 @@ public class HotspotNetworkEntry extends WifiEntry {
     })
     public @interface DeviceType {} // TODO(b/271868642): Add IfThisThanThat lint
 
+    public static final int CONNECTION_STATUS_CONNECTED = 10;
+
     /**
      * If editing this IntDef also edit the definition in:
      * {@link android.net.wifi.sharedconnectivity.app.HotspotNetworkConnectionStatus}
@@ -116,6 +118,7 @@ public class HotspotNetworkEntry extends WifiEntry {
             HotspotNetworkConnectionStatus.CONNECTION_STATUS_ENABLING_HOTSPOT_FAILED,
             HotspotNetworkConnectionStatus.CONNECTION_STATUS_ENABLING_HOTSPOT_TIMEOUT,
             HotspotNetworkConnectionStatus.CONNECTION_STATUS_CONNECT_TO_HOTSPOT_FAILED,
+            CONNECTION_STATUS_CONNECTED,
     })
     public @interface ConnectionStatus {} // TODO(b/271868642): Add IfThisThanThat lint
 
@@ -424,6 +427,17 @@ public class HotspotNetworkEntry extends WifiEntry {
                 mCalledConnect = false;
                 notifyOnUpdated();
                 break;
+            case CONNECTION_STATUS_CONNECTED:
+                mCallbackHandler.post(() -> {
+                    final ConnectCallback connectCallback = mConnectCallback;
+                    if (connectCallback != null) {
+                        connectCallback.onConnectResult(
+                                ConnectCallback.CONNECT_STATUS_SUCCESS);
+                    }
+                });
+                mCalledConnect = false;
+                notifyOnUpdated();
+                break;
             default:
                 // Do nothing
         }
diff --git a/libs/WifiTrackerLib/src/com/android/wifitrackerlib/WifiPickerTracker.java b/libs/WifiTrackerLib/src/com/android/wifitrackerlib/WifiPickerTracker.java
index d36c62594..3ee9a4dad 100644
--- a/libs/WifiTrackerLib/src/com/android/wifitrackerlib/WifiPickerTracker.java
+++ b/libs/WifiTrackerLib/src/com/android/wifitrackerlib/WifiPickerTracker.java
@@ -95,6 +95,9 @@ public class WifiPickerTracker extends BaseWifiTracker {
 
     private static final String TAG = "WifiPickerTracker";
 
+    private static final String EXTRA_KEY_CONNECTION_STATUS_CONNECTED =
+            "connection_status_connected";
+
     private final WifiPickerTrackerCallback mListener;
 
     // Lock object for data returned by the public API
@@ -479,15 +482,25 @@ public class WifiPickerTracker extends BaseWifiTracker {
             updateWifiEntries();
         }
     }
-    @TargetApi(VERSION_CODES.UPSIDE_DOWN_CAKE)
-    @WorkerThread
-    protected void handleHotspotNetworkConnectionStatusChanged(
-            @NonNull HotspotNetworkConnectionStatus status) {
-        mHotspotNetworkEntryCache.stream().filter(
-                entry -> entry.getHotspotNetworkEntryKey().getDeviceId()
-                        == status.getHotspotNetwork().getDeviceId()).forEach(
-                                entry -> entry.onConnectionStatusChanged(status.getStatus()));
-    }
+
+  @TargetApi(VERSION_CODES.UPSIDE_DOWN_CAKE)
+  @WorkerThread
+  protected void handleHotspotNetworkConnectionStatusChanged(
+      @NonNull HotspotNetworkConnectionStatus status) {
+    mHotspotNetworkEntryCache.stream()
+        .filter(
+            entry ->
+                entry.getHotspotNetworkEntryKey().getDeviceId()
+                    == status.getHotspotNetwork().getDeviceId())
+        .forEach(
+            entry -> {
+              if (status.getExtras().getBoolean(EXTRA_KEY_CONNECTION_STATUS_CONNECTED, false)) {
+                entry.onConnectionStatusChanged(HotspotNetworkEntry.CONNECTION_STATUS_CONNECTED);
+              } else {
+                entry.onConnectionStatusChanged(status.getStatus());
+              }
+            });
+  }
 
     @TargetApi(VERSION_CODES.UPSIDE_DOWN_CAKE)
     @WorkerThread
diff --git a/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/HotspotNetworkEntryTest.java b/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/HotspotNetworkEntryTest.java
index 534646066..e4168e489 100644
--- a/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/HotspotNetworkEntryTest.java
+++ b/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/HotspotNetworkEntryTest.java
@@ -652,4 +652,40 @@ public class HotspotNetworkEntryTest {
 
         assertThat(entry.getSummary()).isNotEqualTo("Connecting…");
     }
+
+    @Test
+    public void testOnConnectionStatusChanged_connectedStatus_updatesString() {
+        final HotspotNetworkEntry entry = new HotspotNetworkEntry(
+                mMockInjector, mMockContext, mTestHandler,
+                mMockWifiManager, mMockSharedConnectivityManager, TEST_HOTSPOT_NETWORK_DATA);
+        entry.setListener(mMockListener);
+        entry.onConnectionStatusChanged(
+                HotspotNetworkConnectionStatus.CONNECTION_STATUS_ENABLING_HOTSPOT);
+        mTestLooper.dispatchAll();
+        assertThat(entry.getSummary()).isEqualTo("Connecting…");
+
+        entry.onConnectionStatusChanged(HotspotNetworkEntry.CONNECTION_STATUS_CONNECTED);
+        mTestLooper.dispatchAll();
+
+        assertThat(entry.getSummary()).isNotEqualTo("Connecting…");
+    }
+
+    @Test
+    public void testOnConnectionStatusChanged_connectedStatus_callsCallback() {
+        final HotspotNetworkEntry entry = new HotspotNetworkEntry(
+                mMockInjector, mMockContext, mTestHandler,
+                mMockWifiManager, mMockSharedConnectivityManager, TEST_HOTSPOT_NETWORK_DATA);
+        entry.setListener(mMockListener);
+        entry.connect(mMockConnectCallback);
+        entry.onConnectionStatusChanged(
+                HotspotNetworkConnectionStatus.CONNECTION_STATUS_ENABLING_HOTSPOT);
+        mTestLooper.dispatchAll();
+        verify(mMockConnectCallback, never()).onConnectResult(anyInt());
+
+        entry.onConnectionStatusChanged(HotspotNetworkEntry.CONNECTION_STATUS_CONNECTED);
+        mTestLooper.dispatchAll();
+
+        verify(mMockConnectCallback)
+                .onConnectResult(WifiEntry.ConnectCallback.CONNECT_STATUS_SUCCESS);
+    }
 }
diff --git a/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/WifiPickerTrackerTest.java b/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/WifiPickerTrackerTest.java
index 0b70fed7b..6062fe112 100644
--- a/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/WifiPickerTrackerTest.java
+++ b/libs/WifiTrackerLib/tests/src/com/android/wifitrackerlib/WifiPickerTrackerTest.java
@@ -71,6 +71,7 @@ import android.net.wifi.sharedconnectivity.app.KnownNetworkConnectionStatus;
 import android.net.wifi.sharedconnectivity.app.NetworkProviderInfo;
 import android.net.wifi.sharedconnectivity.app.SharedConnectivityClientCallback;
 import android.net.wifi.sharedconnectivity.app.SharedConnectivityManager;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.test.TestLooper;
 import android.telephony.SubscriptionManager;
@@ -2931,4 +2932,46 @@ public class WifiPickerTrackerTest {
         assertThat(((HotspotNetworkEntry) wifiPickerTracker.getWifiEntries().get(2))
                 .getHotspotNetworkEntryKey().getDeviceId()).isEqualTo(1);
     }
+
+    @Test
+    public void testHotspotNetworks_onHotspotNetworkConnectionStatusChanged_connectedExtra() {
+        final HotspotNetwork testHotspotNetwork =
+            new HotspotNetwork.Builder()
+                .setDeviceId(1)
+                .setNetworkProviderInfo(
+                    new NetworkProviderInfo.Builder("My Phone", "Pixel 7")
+                        .setDeviceType(NetworkProviderInfo.DEVICE_TYPE_PHONE)
+                        .setBatteryPercentage(100)
+                        .setConnectionStrength(3)
+                        .build())
+                .setHostNetworkType(HotspotNetwork.NETWORK_TYPE_CELLULAR)
+                .setNetworkName("Google Fi")
+                .build();
+        when(mMockSharedConnectivityManager.getHotspotNetworks())
+            .thenReturn(Collections.singletonList(testHotspotNetwork));
+        final WifiPickerTracker wifiPickerTracker = createTestWifiPickerTracker();
+        wifiPickerTracker.onStart();
+        mTestLooper.dispatchAll();
+        verify(mMockSharedConnectivityManager)
+            .registerCallback(any(), mSharedConnectivityCallbackCaptor.capture());
+        mSharedConnectivityCallbackCaptor.getValue().onServiceConnected();
+        mTestLooper.dispatchAll();
+
+        final WifiEntry.ConnectCallback connectCallback = mock(WifiEntry.ConnectCallback.class);
+        wifiPickerTracker.getWifiEntries().get(0).connect(connectCallback);
+
+        Bundle extras = new Bundle();
+        extras.putBoolean("connection_status_connected", true);
+        mSharedConnectivityCallbackCaptor
+            .getValue()
+            .onHotspotNetworkConnectionStatusChanged(
+                new HotspotNetworkConnectionStatus.Builder()
+                    .setStatus(HotspotNetworkConnectionStatus.CONNECTION_STATUS_UNKNOWN)
+                    .setExtras(extras)
+                    .setHotspotNetwork(testHotspotNetwork)
+                    .build());
+        mTestLooper.dispatchAll();
+
+        verify(connectCallback).onConnectResult(WifiEntry.ConnectCallback.CONNECT_STATUS_SUCCESS);
+    }
 }