Fix mDNS offload mgr test suite

Also move ConnectivityManager callbacks onto the handler thread.

Test: atest MdnsOffloadManagerTest
Bug: 232507295
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:4335ac2ea2059a4249e274b6a6ff030870631191)
Merged-In: I46b809839f827ebe765a16cd1fca9ce63a62d1fe
Change-Id: I46b809839f827ebe765a16cd1fca9ce63a62d1fe

3 files changed, 54 insertions, 18 deletions

diff --git a/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerService.java b/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerService.java
index b853185..c554cf8 100644
--- a/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerService.java
+++ b/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerService.java
@@ -44,13 +44,14 @@ import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 import androidx.annotation.WorkerThread;
 
+import com.android.tv.mdnsoffloadmanager.util.WakeLockWrapper;
+
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -76,7 +77,7 @@ public class MdnsOffloadManagerService extends Service {
     private OffloadWriter mOffloadWriter;
     private ConnectivityManager mConnectivityManager;
     private PackageManager mPackageManager;
-    private PowerManager.WakeLock mWakeLock;
+    private WakeLockWrapper mWakeLock;
 
     public MdnsOffloadManagerService() {
         this(new Injector());
@@ -121,9 +122,10 @@ public class MdnsOffloadManagerService extends Service {
             return mContext.getSystemService(PowerManager.class).getLowPowerStandbyPolicy();
         }
 
-        PowerManager.WakeLock newWakeLock() {
-            return mContext.getSystemService(PowerManager.class).newWakeLock(
-                    PowerManager.PARTIAL_WAKE_LOCK, TAG);
+        WakeLockWrapper newWakeLock() {
+            return new WakeLockWrapper(
+                    mContext.getSystemService(PowerManager.class)
+                            .newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG));
         }
 
         PackageManager getPackageManager() {
@@ -433,16 +435,15 @@ public class MdnsOffloadManagerService extends Service {
     }
 
     private class ConnectivityManagerNetworkCallback extends ConnectivityManager.NetworkCallback {
-        private Map<Network, LinkProperties> mLinkPropertiesMap = new ConcurrentHashMap<>();
+        private final Map<Network, LinkProperties> mLinkProperties = new HashMap<>();
 
         @Override
-        public void onLinkPropertiesChanged(Network network,
-                LinkProperties linkProperties) {
+        public void onLinkPropertiesChanged(Network network, LinkProperties linkProperties) {
             // We only want to know the interface name of a network. This method is
             // called right after onAvailable() or any other important change during the lifecycle
             // of the network.
-            LinkProperties previousProperties = mLinkPropertiesMap.put(network, linkProperties);
             mHandler.post(() -> {
+                LinkProperties previousProperties = mLinkProperties.put(network, linkProperties);
                 if (previousProperties != null &&
                         !previousProperties.getInterfaceName().equals(
                                 linkProperties.getInterfaceName())) {
@@ -459,20 +460,19 @@ public class MdnsOffloadManagerService extends Service {
                         getInterfaceOffloadManager(linkProperties.getInterfaceName());
                 offloadManager.onNetworkAvailable();
             });
-
         }
 
         @Override
         public void onLost(@NonNull Network network) {
-            // Network object is guaranteed to match a network object from a previous
-            // onLinkPropertiesChanged() so the LinkProperties must be available to retrieve
-            // the associated iface.
-            LinkProperties previousProperties = mLinkPropertiesMap.remove(network);
-            if (previousProperties == null){
-                Log.w(TAG,"Network "+ network + " lost before being available.");
-                return;
-            }
             mHandler.post(() -> {
+                // Network object is guaranteed to match a network object from a previous
+                // onLinkPropertiesChanged() so the LinkProperties must be available to retrieve
+                // the associated iface.
+                LinkProperties previousProperties = mLinkProperties.remove(network);
+                if (previousProperties == null){
+                    Log.w(TAG,"Network "+ network + " lost before being available.");
+                    return;
+                }
                 InterfaceOffloadManager offloadManager =
                         getInterfaceOffloadManager(previousProperties.getInterfaceName());
                 offloadManager.onNetworkLost();
diff --git a/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/util/WakeLockWrapper.java b/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/util/WakeLockWrapper.java
new file mode 100644
index 0000000..150d801
--- /dev/null
+++ b/MdnsOffloadManagerService/src/com/android/tv/mdnsoffloadmanager/util/WakeLockWrapper.java
@@ -0,0 +1,28 @@
+package com.android.tv.mdnsoffloadmanager.util;
+
+import android.os.PowerManager;
+
+/**
+ * Wrapper around {@link android.os.PowerManager.WakeLock} for testing purposes.
+ */
+public class WakeLockWrapper {
+    private final PowerManager.WakeLock mLock;
+
+    public WakeLockWrapper(PowerManager.WakeLock lock) {
+        this.mLock = lock;
+    }
+
+    /**
+     * @see PowerManager.WakeLock#acquire()
+     */
+    public void acquire(long timeout) {
+        mLock.acquire(timeout);
+    }
+
+    /**
+     * @see PowerManager.WakeLock#release()
+     */
+    public void release() {
+        mLock.release();
+    }
+}
diff --git a/MdnsOffloadManagerService/tests/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerTest.java b/MdnsOffloadManagerService/tests/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerTest.java
index fe03990..d5f2073 100644
--- a/MdnsOffloadManagerService/tests/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerTest.java
+++ b/MdnsOffloadManagerService/tests/src/com/android/tv/mdnsoffloadmanager/MdnsOffloadManagerTest.java
@@ -69,6 +69,7 @@ import android.util.Log;
 import androidx.test.filters.SmallTest;
 
 import com.android.tv.mdnsoffloadmanager.MdnsOffloadManagerService.Injector;
+import com.android.tv.mdnsoffloadmanager.util.WakeLockWrapper;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -120,6 +121,8 @@ public class MdnsOffloadManagerTest {
     Network mNetwork1;
     @Mock
     PackageManager mPackageManager;
+    @Mock
+    WakeLockWrapper mWakeLock;
     @Spy
     FakeMdnsOffloadService mVendorService = new FakeMdnsOffloadService();
     @Captor
@@ -185,6 +188,11 @@ public class MdnsOffloadManagerTest {
             }
 
             @Override
+            WakeLockWrapper newWakeLock() {
+                return mWakeLock;
+            }
+
+            @Override
             PackageManager getPackageManager() {
                 return mPackageManager;
             }