Merge "Restore connected state in Wifi picker." into main

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);
+    }
 }