Merge "Refine NetworkCallbackTest" into rvc-dev
diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
index a820ae5..5aafdf0 100644
--- a/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
@@ -25,4 +25,5 @@
     String getRestrictBackgroundStatus();
     void sendNotification(int notificationId, String notificationType);
     void registerNetworkCallback(in INetworkCallback cb);
+    void unregisterNetworkCallback();
 }
diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkCallback.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkCallback.aidl
index 740ec26..2048bab 100644
--- a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkCallback.aidl
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkCallback.aidl
@@ -17,9 +17,11 @@
 package com.android.cts.net.hostside;
 
 import android.net.Network;
+import android.net.NetworkCapabilities;
 
 interface INetworkCallback {
     void onBlockedStatusChanged(in Network network, boolean blocked);
     void onAvailable(in Network network);
     void onLost(in Network network);
+    void onCapabilitiesChanged(in Network network, in NetworkCapabilities cap);
 }
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
index 1db0417..2072db3 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
@@ -700,6 +700,10 @@
         mServiceClient.registerNetworkCallback(cb);
     }
 
+    protected void unregisterNetworkCallback() throws Exception {
+        mServiceClient.unregisterNetworkCallback();
+    }
+
     /**
      * Registers a {@link NotificationListenerService} implementation that will execute the
      * notification actions right after the notification is sent.
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
index 3ee7b99..6546e26 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
@@ -100,4 +100,8 @@
     public void registerNetworkCallback(INetworkCallback cb) throws RemoteException {
         mService.registerNetworkCallback(cb);
     }
+
+    public void unregisterNetworkCallback() throws RemoteException {
+        mService.unregisterNetworkCallback();
+    }
 }
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkCallbackTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkCallbackTest.java
index ed397b9..ec884d0 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkCallbackTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkCallbackTest.java
@@ -16,18 +16,23 @@
 
 package com.android.cts.net.hostside;
 
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.setRestrictBackground;
 import static com.android.cts.net.hostside.Property.BATTERY_SAVER_MODE;
 import static com.android.cts.net.hostside.Property.DATA_SAVER_MODE;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import android.net.Network;
+import android.net.NetworkCapabilities;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Objects;
@@ -35,15 +40,18 @@
 import java.util.concurrent.TimeUnit;
 
 public class NetworkCallbackTest extends AbstractRestrictBackgroundNetworkTestCase {
-
     private Network mNetwork;
     private final TestNetworkCallback mTestNetworkCallback = new TestNetworkCallback();
+    @Rule
+    public final MeterednessConfigurationRule mMeterednessConfiguration
+            = new MeterednessConfigurationRule();
 
     enum CallbackState {
         NONE,
         AVAILABLE,
         LOST,
-        BLOCKED_STATUS
+        BLOCKED_STATUS,
+        CAPABILITIES
     }
 
     private static class CallbackInfo {
@@ -75,7 +83,7 @@
     }
 
     private class TestNetworkCallback extends INetworkCallback.Stub {
-        private static final int TEST_CALLBACK_TIMEOUT_MS = 200;
+        private static final int TEST_CALLBACK_TIMEOUT_MS = 5_000;
 
         private final LinkedBlockingQueue<CallbackInfo> mCallbacks = new LinkedBlockingQueue<>();
 
@@ -117,12 +125,21 @@
             setLastCallback(CallbackState.BLOCKED_STATUS, network, blocked);
         }
 
+        @Override
+        public void onCapabilitiesChanged(Network network, NetworkCapabilities cap) {
+            setLastCallback(CallbackState.CAPABILITIES, network, cap);
+        }
+
         public void expectLostCallback(Network expectedNetwork) {
             expectCallback(CallbackState.LOST, expectedNetwork, null);
         }
 
-        public void expectAvailableCallback(Network expectedNetwork) {
-            expectCallback(CallbackState.AVAILABLE, expectedNetwork, null);
+        public Network expectAvailableCallbackAndGetNetwork() {
+            final CallbackInfo cb = nextCallback(TEST_CALLBACK_TIMEOUT_MS);
+            if (cb.state != CallbackState.AVAILABLE) {
+                fail("Network is not available");
+            }
+            return cb.network;
         }
 
         public void expectBlockedStatusCallback(Network expectedNetwork, boolean expectBlocked) {
@@ -130,15 +147,28 @@
                     expectBlocked);
         }
 
-        void assertNoCallback() {
-            CallbackInfo cb = null;
-            try {
-                cb = mCallbacks.poll(TEST_CALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
-            } catch (InterruptedException e) {
-                // Expected.
-            }
-            if (cb != null) {
-                assertNull("Unexpected callback: " + cb, cb);
+        public void waitBlockedStatusCallback(Network expectedNetwork, boolean expectBlocked) {
+            final long deadline = System.currentTimeMillis() + TEST_CALLBACK_TIMEOUT_MS;
+            do {
+                final CallbackInfo cb = nextCallback((int) (deadline - System.currentTimeMillis()));
+                if (cb.state == CallbackState.BLOCKED_STATUS) {
+                    assertEquals(expectBlocked, (Boolean) cb.arg);
+                    return;
+                }
+            } while (System.currentTimeMillis() <= deadline);
+            fail("Didn't receive onBlockedStatusChanged()");
+        }
+
+        public void expectCapabilitiesCallback(Network expectedNetwork, boolean hasCapability,
+                int capability) {
+            final CallbackInfo cb = nextCallback(TEST_CALLBACK_TIMEOUT_MS);
+            final NetworkCapabilities cap = (NetworkCapabilities) cb.arg;
+            assertEquals(expectedNetwork, cb.network);
+            assertEquals(CallbackState.CAPABILITIES, cb.state);
+            if (hasCapability) {
+                assertTrue(cap.hasCapability(capability));
+            } else {
+                assertFalse(cap.hasCapability(capability));
             }
         }
     }
@@ -147,13 +177,28 @@
     public void setUp() throws Exception {
         super.setUp();
 
-        mNetwork = mCm.getActiveNetwork();
-
         registerBroadcastReceiver();
 
         removeRestrictBackgroundWhitelist(mUid);
         removeRestrictBackgroundBlacklist(mUid);
         assertRestrictBackgroundChangedReceived(0);
+
+        // Initial state
+        setBatterySaverMode(false);
+        setRestrictBackground(false);
+
+        // Make wifi a metered network.
+        mMeterednessConfiguration.configureNetworkMeteredness(true);
+
+        // Register callback
+        registerNetworkCallback((INetworkCallback.Stub) mTestNetworkCallback);
+        // Once the wifi is marked as metered, the wifi will reconnect. Wait for onAvailable()
+        // callback to ensure wifi is connected before the test and store the default network.
+        mNetwork = mTestNetworkCallback.expectAvailableCallbackAndGetNetwork();
+        // Check that the network is metered.
+        mTestNetworkCallback.expectCapabilitiesCallback(mNetwork, false /* hasCapability */,
+                NET_CAPABILITY_NOT_METERED);
+        mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
     }
 
     @After
@@ -162,102 +207,80 @@
 
         setRestrictBackground(false);
         setBatterySaverMode(false);
+        unregisterNetworkCallback();
     }
 
     @RequiredProperties({DATA_SAVER_MODE})
     @Test
     public void testOnBlockedStatusChanged_dataSaver() throws Exception {
-        // Initial state
-        setBatterySaverMode(false);
-        setRestrictBackground(false);
-
-        final MeterednessConfigurationRule meterednessConfiguration
-                = new MeterednessConfigurationRule();
-        meterednessConfiguration.configureNetworkMeteredness(true);
         try {
-            // Register callback
-            registerNetworkCallback((INetworkCallback.Stub) mTestNetworkCallback);
-            mTestNetworkCallback.expectAvailableCallback(mNetwork);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
-
             // Enable restrict background
             setRestrictBackground(true);
             assertBackgroundNetworkAccess(false);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, true);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, true);
 
             // Add to whitelist
             addRestrictBackgroundWhitelist(mUid);
             assertBackgroundNetworkAccess(true);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, false);
 
             // Remove from whitelist
             removeRestrictBackgroundWhitelist(mUid);
             assertBackgroundNetworkAccess(false);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, true);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, true);
         } finally {
-            meterednessConfiguration.resetNetworkMeteredness();
+            mMeterednessConfiguration.resetNetworkMeteredness();
         }
 
         // Set to non-metered network
-        meterednessConfiguration.configureNetworkMeteredness(false);
+        mMeterednessConfiguration.configureNetworkMeteredness(false);
+        mTestNetworkCallback.expectCapabilitiesCallback(mNetwork, true /* hasCapability */,
+                NET_CAPABILITY_NOT_METERED);
         try {
             assertBackgroundNetworkAccess(true);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, false);
 
             // Disable restrict background, should not trigger callback
             setRestrictBackground(false);
             assertBackgroundNetworkAccess(true);
-            mTestNetworkCallback.assertNoCallback();
         } finally {
-            meterednessConfiguration.resetNetworkMeteredness();
+            mMeterednessConfiguration.resetNetworkMeteredness();
         }
     }
 
     @RequiredProperties({BATTERY_SAVER_MODE})
     @Test
     public void testOnBlockedStatusChanged_powerSaver() throws Exception {
-        // Set initial state.
-        setBatterySaverMode(false);
-        setRestrictBackground(false);
-
-        final MeterednessConfigurationRule meterednessConfiguration
-                = new MeterednessConfigurationRule();
-        meterednessConfiguration.configureNetworkMeteredness(true);
         try {
-            // Register callback
-            registerNetworkCallback((INetworkCallback.Stub) mTestNetworkCallback);
-            mTestNetworkCallback.expectAvailableCallback(mNetwork);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
-
             // Enable Power Saver
             setBatterySaverMode(true);
             assertBackgroundNetworkAccess(false);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, true);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, true);
 
             // Disable Power Saver
             setBatterySaverMode(false);
             assertBackgroundNetworkAccess(true);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, false);
         } finally {
-            meterednessConfiguration.resetNetworkMeteredness();
+            mMeterednessConfiguration.resetNetworkMeteredness();
         }
 
         // Set to non-metered network
-        meterednessConfiguration.configureNetworkMeteredness(false);
+        mMeterednessConfiguration.configureNetworkMeteredness(false);
+        mTestNetworkCallback.expectCapabilitiesCallback(mNetwork, true /* hasCapability */,
+                NET_CAPABILITY_NOT_METERED);
         try {
-            mTestNetworkCallback.assertNoCallback();
-
             // Enable Power Saver
             setBatterySaverMode(true);
             assertBackgroundNetworkAccess(false);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, true);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, true);
 
             // Disable Power Saver
             setBatterySaverMode(false);
             assertBackgroundNetworkAccess(true);
-            mTestNetworkCallback.expectBlockedStatusCallback(mNetwork, false);
+            mTestNetworkCallback.waitBlockedStatusCallback(mNetwork, false);
         } finally {
-            meterednessConfiguration.resetNetworkMeteredness();
+            mMeterednessConfiguration.resetNetworkMeteredness();
         }
     }
 
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
index ec536af..590e17e 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
@@ -127,6 +127,16 @@
                         unregisterNetworkCallback();
                     }
                 }
+
+                @Override
+                public void onCapabilitiesChanged(Network network, NetworkCapabilities cap) {
+                    try {
+                        cb.onCapabilitiesChanged(network, cap);
+                    } catch (RemoteException e) {
+                        Log.d(TAG, "Cannot send onCapabilitiesChanged: " + e);
+                        unregisterNetworkCallback();
+                    }
+                }
             };
             mCm.registerNetworkCallback(makeWifiNetworkRequest(), mNetworkCallback);
             try {
@@ -135,17 +145,21 @@
                 unregisterNetworkCallback();
             }
         }
-      };
 
-    private void unregisterNetworkCallback() {
-        Log.d(TAG, "unregistering network callback");
-        mCm.unregisterNetworkCallback(mNetworkCallback);
-        mNetworkCallback = null;
-    }
+        @Override
+        public void unregisterNetworkCallback() {
+            Log.d(TAG, "unregistering network callback");
+            if (mNetworkCallback != null) {
+                mCm.unregisterNetworkCallback(mNetworkCallback);
+                mNetworkCallback = null;
+            }
+        }
+      };
 
     private NetworkRequest makeWifiNetworkRequest() {
         return new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
                 .build();
     }