Store priority class as final in UnderlyingNetworkRecord

This commit refactors UnderlyingNetworkRecord to store
its priority class as a final field and calculate it
during construction. This commit improves the code
maintainabiliy in the network selection logic.

Bug: 245618320
Test: FrameworksVcnTests, CtsVcnTestCases
Change-Id: I346323cb18de935626b2dc6343be41c21fc7f055

6 files changed, 182 insertions, 147 deletions

diff --git a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
index d474c5d33a93..3acbb5e17e13 100644
--- a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
+++ b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
@@ -356,7 +356,7 @@ public class UnderlyingNetworkController {
             if (!allNetworkPriorities.isEmpty()) {
                 allNetworkPriorities += ", ";
             }
-            allNetworkPriorities += record.network + ": " + record.getPriorityClass();
+            allNetworkPriorities += record.network + ": " + record.priorityClass;
         }
         logInfo(
                 "Selected network changed to "
@@ -393,19 +393,20 @@ public class UnderlyingNetworkController {
 
         private TreeSet<UnderlyingNetworkRecord> getSortedUnderlyingNetworks() {
             TreeSet<UnderlyingNetworkRecord> sorted =
-                    new TreeSet<>(
-                            UnderlyingNetworkRecord.getComparator(
+                    new TreeSet<>(UnderlyingNetworkRecord.getComparator());
+
+            for (UnderlyingNetworkRecord.Builder builder :
+                    mUnderlyingNetworkRecordBuilders.values()) {
+                if (builder.isValid()) {
+                    final UnderlyingNetworkRecord record =
+                            builder.build(
                                     mVcnContext,
                                     mConnectionConfig.getVcnUnderlyingNetworkPriorities(),
                                     mSubscriptionGroup,
                                     mLastSnapshot,
                                     mCurrentRecord,
-                                    mCarrierConfig));
-
-            for (UnderlyingNetworkRecord.Builder builder :
-                    mUnderlyingNetworkRecordBuilders.values()) {
-                if (builder.isValid()) {
-                    sorted.add(builder.build());
+                                    mCarrierConfig);
+                    sorted.add(record);
                 }
             }
 
diff --git a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkRecord.java b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkRecord.java
index 319680e0b01c..aea9f4d2dbae 100644
--- a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkRecord.java
+++ b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkRecord.java
@@ -42,53 +42,58 @@ import java.util.Objects;
  * @hide
  */
 public class UnderlyingNetworkRecord {
-    private static final int PRIORITY_CLASS_INVALID = Integer.MAX_VALUE;
-
     @NonNull public final Network network;
     @NonNull public final NetworkCapabilities networkCapabilities;
     @NonNull public final LinkProperties linkProperties;
     public final boolean isBlocked;
-
-    private int mPriorityClass = PRIORITY_CLASS_INVALID;
+    public final boolean isSelected;
+    public final int priorityClass;
 
     @VisibleForTesting(visibility = Visibility.PRIVATE)
     public UnderlyingNetworkRecord(
             @NonNull Network network,
             @NonNull NetworkCapabilities networkCapabilities,
             @NonNull LinkProperties linkProperties,
-            boolean isBlocked) {
-        this.network = network;
-        this.networkCapabilities = networkCapabilities;
-        this.linkProperties = linkProperties;
-        this.isBlocked = isBlocked;
-    }
-
-    private int getOrCalculatePriorityClass(
+            boolean isBlocked,
             VcnContext vcnContext,
             List<VcnUnderlyingNetworkTemplate> underlyingNetworkTemplates,
             ParcelUuid subscriptionGroup,
             TelephonySubscriptionSnapshot snapshot,
             UnderlyingNetworkRecord currentlySelected,
             PersistableBundleWrapper carrierConfig) {
-        // Never changes after the underlying network record is created.
-        if (mPriorityClass == PRIORITY_CLASS_INVALID) {
-            mPriorityClass =
-                    NetworkPriorityClassifier.calculatePriorityClass(
-                            vcnContext,
-                            this,
-                            underlyingNetworkTemplates,
-                            subscriptionGroup,
-                            snapshot,
-                            currentlySelected,
-                            carrierConfig);
-        }
+        this.network = network;
+        this.networkCapabilities = networkCapabilities;
+        this.linkProperties = linkProperties;
+        this.isBlocked = isBlocked;
 
-        return mPriorityClass;
+        this.isSelected = isSelected(this.network, currentlySelected);
+
+        priorityClass =
+                NetworkPriorityClassifier.calculatePriorityClass(
+                        vcnContext,
+                        this,
+                        underlyingNetworkTemplates,
+                        subscriptionGroup,
+                        snapshot,
+                        currentlySelected,
+                        carrierConfig);
     }
 
-    // Used in UnderlyingNetworkController
-    int getPriorityClass() {
-        return mPriorityClass;
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    public UnderlyingNetworkRecord(
+            @NonNull Network network,
+            @NonNull NetworkCapabilities networkCapabilities,
+            @NonNull LinkProperties linkProperties,
+            boolean isBlocked,
+            boolean isSelected,
+            int priorityClass) {
+        this.network = network;
+        this.networkCapabilities = networkCapabilities;
+        this.linkProperties = linkProperties;
+        this.isBlocked = isBlocked;
+        this.isSelected = isSelected;
+
+        this.priorityClass = priorityClass;
     }
 
     @Override
@@ -108,40 +113,32 @@ public class UnderlyingNetworkRecord {
         return Objects.hash(network, networkCapabilities, linkProperties, isBlocked);
     }
 
-    static Comparator<UnderlyingNetworkRecord> getComparator(
-            VcnContext vcnContext,
-            List<VcnUnderlyingNetworkTemplate> underlyingNetworkTemplates,
-            ParcelUuid subscriptionGroup,
-            TelephonySubscriptionSnapshot snapshot,
-            UnderlyingNetworkRecord currentlySelected,
-            PersistableBundleWrapper carrierConfig) {
+    /** Returns if two records are equal including their priority classes. */
+    public static boolean isEqualIncludingPriorities(
+            UnderlyingNetworkRecord left, UnderlyingNetworkRecord right) {
+        if (left != null && right != null) {
+            return left.equals(right)
+                    && left.isSelected == right.isSelected
+                    && left.priorityClass == right.priorityClass;
+        }
+
+        return left == right;
+    }
+
+    static Comparator<UnderlyingNetworkRecord> getComparator() {
         return (left, right) -> {
-            final int leftIndex =
-                    left.getOrCalculatePriorityClass(
-                            vcnContext,
-                            underlyingNetworkTemplates,
-                            subscriptionGroup,
-                            snapshot,
-                            currentlySelected,
-                            carrierConfig);
-            final int rightIndex =
-                    right.getOrCalculatePriorityClass(
-                            vcnContext,
-                            underlyingNetworkTemplates,
-                            subscriptionGroup,
-                            snapshot,
-                            currentlySelected,
-                            carrierConfig);
+            final int leftIndex = left.priorityClass;
+            final int rightIndex = right.priorityClass;
 
             // In the case of networks in the same priority class, prioritize based on other
             // criteria (eg. actively selected network, link metrics, etc)
             if (leftIndex == rightIndex) {
                 // TODO: Improve the strategy of network selection when both UnderlyingNetworkRecord
                 // fall into the same priority class.
-                if (isSelected(left, currentlySelected)) {
+                if (left.isSelected) {
                     return -1;
                 }
-                if (isSelected(left, currentlySelected)) {
+                if (right.isSelected) {
                     return 1;
                 }
             }
@@ -150,11 +147,11 @@ public class UnderlyingNetworkRecord {
     }
 
     private static boolean isSelected(
-            UnderlyingNetworkRecord recordToCheck, UnderlyingNetworkRecord currentlySelected) {
+            Network networkToCheck, UnderlyingNetworkRecord currentlySelected) {
         if (currentlySelected == null) {
             return false;
         }
-        if (currentlySelected.network == recordToCheck.network) {
+        if (currentlySelected.network.equals(networkToCheck)) {
             return true;
         }
         return false;
@@ -172,16 +169,8 @@ public class UnderlyingNetworkRecord {
         pw.println("UnderlyingNetworkRecord:");
         pw.increaseIndent();
 
-        final int priorityIndex =
-                getOrCalculatePriorityClass(
-                        vcnContext,
-                        underlyingNetworkTemplates,
-                        subscriptionGroup,
-                        snapshot,
-                        currentlySelected,
-                        carrierConfig);
-
-        pw.println("Priority index: " + priorityIndex);
+        pw.println("priorityClass: " + priorityClass);
+        pw.println("isSelected: " + isSelected);
         pw.println("mNetwork: " + network);
         pw.println("mNetworkCapabilities: " + networkCapabilities);
         pw.println("mLinkProperties: " + linkProperties);
@@ -198,8 +187,6 @@ public class UnderlyingNetworkRecord {
         boolean mIsBlocked;
         boolean mWasIsBlockedSet;
 
-        @Nullable private UnderlyingNetworkRecord mCached;
-
         Builder(@NonNull Network network) {
             mNetwork = network;
         }
@@ -211,7 +198,6 @@ public class UnderlyingNetworkRecord {
 
         void setNetworkCapabilities(@NonNull NetworkCapabilities networkCapabilities) {
             mNetworkCapabilities = networkCapabilities;
-            mCached = null;
         }
 
         @Nullable
@@ -221,32 +207,40 @@ public class UnderlyingNetworkRecord {
 
         void setLinkProperties(@NonNull LinkProperties linkProperties) {
             mLinkProperties = linkProperties;
-            mCached = null;
         }
 
         void setIsBlocked(boolean isBlocked) {
             mIsBlocked = isBlocked;
             mWasIsBlockedSet = true;
-            mCached = null;
         }
 
         boolean isValid() {
             return mNetworkCapabilities != null && mLinkProperties != null && mWasIsBlockedSet;
         }
 
-        UnderlyingNetworkRecord build() {
+        UnderlyingNetworkRecord build(
+                VcnContext vcnContext,
+                List<VcnUnderlyingNetworkTemplate> underlyingNetworkTemplates,
+                ParcelUuid subscriptionGroup,
+                TelephonySubscriptionSnapshot snapshot,
+                UnderlyingNetworkRecord currentlySelected,
+                PersistableBundleWrapper carrierConfig) {
             if (!isValid()) {
                 throw new IllegalArgumentException(
                         "Called build before UnderlyingNetworkRecord was valid");
             }
 
-            if (mCached == null) {
-                mCached =
-                        new UnderlyingNetworkRecord(
-                                mNetwork, mNetworkCapabilities, mLinkProperties, mIsBlocked);
-            }
-
-            return mCached;
+            return new UnderlyingNetworkRecord(
+                    mNetwork,
+                    mNetworkCapabilities,
+                    mLinkProperties,
+                    mIsBlocked,
+                    vcnContext,
+                    underlyingNetworkTemplates,
+                    subscriptionGroup,
+                    snapshot,
+                    currentlySelected,
+                    carrierConfig);
         }
     }
 }
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
index a4ee2de9f433..692c8a8f0898 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
@@ -143,9 +143,9 @@ public class VcnGatewayConnectionTest extends VcnGatewayConnectionTestBase {
         capBuilder.setLinkDownstreamBandwidthKbps(TEST_DOWNSTREAM_BANDWIDTH);
         capBuilder.setAdministratorUids(new int[] {TEST_UID});
         final Network underlyingNetwork = mock(Network.class, CALLS_REAL_METHODS);
-        UnderlyingNetworkRecord record = new UnderlyingNetworkRecord(
-                underlyingNetwork,
-                capBuilder.build(), new LinkProperties(), false);
+        UnderlyingNetworkRecord record =
+                getTestNetworkRecord(
+                        underlyingNetwork, capBuilder.build(), new LinkProperties(), false);
         final NetworkCapabilities vcnCaps =
                 VcnGatewayConnection.buildNetworkCapabilities(
                         VcnGatewayConnectionConfigTest.buildTestConfig(),
@@ -211,7 +211,7 @@ public class VcnGatewayConnectionTest extends VcnGatewayConnectionTestBase {
         doReturn(TEST_DNS_ADDRESSES).when(childSessionConfig).getInternalDnsServers();
 
         UnderlyingNetworkRecord record =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mock(Network.class, CALLS_REAL_METHODS),
                         new NetworkCapabilities.Builder().build(),
                         underlyingLp,
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
index 7bafd243799f..0a18c3d1990d 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
@@ -109,9 +109,23 @@ public class VcnGatewayConnectionTestBase {
     protected static final long ELAPSED_REAL_TIME = 123456789L;
     protected static final String TEST_IPSEC_TUNNEL_IFACE = "IPSEC_IFACE";
 
+    protected static UnderlyingNetworkRecord getTestNetworkRecord(
+            Network network,
+            NetworkCapabilities networkCapabilities,
+            LinkProperties linkProperties,
+            boolean isBlocked) {
+        return new UnderlyingNetworkRecord(
+                network,
+                networkCapabilities,
+                linkProperties,
+                isBlocked,
+                false /* isSelected */,
+                0 /* priorityClass */);
+    }
+
     protected static final String TEST_TCP_BUFFER_SIZES_1 = "1,2,3,4";
     protected static final UnderlyingNetworkRecord TEST_UNDERLYING_NETWORK_RECORD_1 =
-            new UnderlyingNetworkRecord(
+            getTestNetworkRecord(
                     mock(Network.class, CALLS_REAL_METHODS),
                     new NetworkCapabilities(),
                     new LinkProperties(),
@@ -124,7 +138,7 @@ public class VcnGatewayConnectionTestBase {
 
     protected static final String TEST_TCP_BUFFER_SIZES_2 = "2,3,4,5";
     protected static final UnderlyingNetworkRecord TEST_UNDERLYING_NETWORK_RECORD_2 =
-            new UnderlyingNetworkRecord(
+            getTestNetworkRecord(
                     mock(Network.class, CALLS_REAL_METHODS),
                     new NetworkCapabilities(),
                     new LinkProperties(),
diff --git a/tests/vcn/java/com/android/server/vcn/routeselection/NetworkPriorityClassifierTest.java b/tests/vcn/java/com/android/server/vcn/routeselection/NetworkPriorityClassifierTest.java
index b0d68952c39d..bb7da5604e16 100644
--- a/tests/vcn/java/com/android/server/vcn/routeselection/NetworkPriorityClassifierTest.java
+++ b/tests/vcn/java/com/android/server/vcn/routeselection/NetworkPriorityClassifierTest.java
@@ -25,7 +25,6 @@ import static android.net.vcn.VcnUnderlyingNetworkTemplateTestBase.TEST_MIN_EXIT
 
 import static com.android.server.vcn.VcnTestUtils.setupSystemService;
 import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.PRIORITY_ANY;
-import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.calculatePriorityClass;
 import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.checkMatchesCellPriorityRule;
 import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.checkMatchesPriorityRule;
 import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.checkMatchesWifiPriorityRule;
@@ -48,6 +47,7 @@ import android.net.TelephonyNetworkSpecifier;
 import android.net.vcn.VcnCellUnderlyingNetworkTemplate;
 import android.net.vcn.VcnGatewayConnectionConfig;
 import android.net.vcn.VcnManager;
+import android.net.vcn.VcnUnderlyingNetworkTemplate;
 import android.net.vcn.VcnWifiUnderlyingNetworkTemplate;
 import android.os.ParcelUuid;
 import android.os.PersistableBundle;
@@ -64,6 +64,7 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.List;
 import java.util.Set;
 import java.util.UUID;
 
@@ -135,25 +136,35 @@ public class NetworkPriorityClassifierTest {
                                 false /* isInTestMode */));
         doNothing().when(mVcnContext).ensureRunningOnLooperThread();
 
-        mWifiNetworkRecord =
-                new UnderlyingNetworkRecord(
-                        mNetwork,
-                        WIFI_NETWORK_CAPABILITIES,
-                        LINK_PROPERTIES,
-                        false /* isBlocked */);
-
-        mCellNetworkRecord =
-                new UnderlyingNetworkRecord(
-                        mNetwork,
-                        CELL_NETWORK_CAPABILITIES,
-                        LINK_PROPERTIES,
-                        false /* isBlocked */);
-
         setupSystemService(
                 mockContext, mTelephonyManager, Context.TELEPHONY_SERVICE, TelephonyManager.class);
         when(mTelephonyManager.createForSubscriptionId(SUB_ID)).thenReturn(mTelephonyManager);
         when(mTelephonyManager.getNetworkOperator()).thenReturn(PLMN_ID);
         when(mTelephonyManager.getSimSpecificCarrierId()).thenReturn(CARRIER_ID);
+
+        mWifiNetworkRecord =
+                getTestNetworkRecord(
+                        WIFI_NETWORK_CAPABILITIES,
+                        VcnGatewayConnectionConfig.DEFAULT_UNDERLYING_NETWORK_TEMPLATES);
+        mCellNetworkRecord =
+                getTestNetworkRecord(
+                        CELL_NETWORK_CAPABILITIES,
+                        VcnGatewayConnectionConfig.DEFAULT_UNDERLYING_NETWORK_TEMPLATES);
+    }
+
+    private UnderlyingNetworkRecord getTestNetworkRecord(
+            NetworkCapabilities nc, List<VcnUnderlyingNetworkTemplate> underlyingNetworkTemplates) {
+        return new UnderlyingNetworkRecord(
+                mNetwork,
+                nc,
+                LINK_PROPERTIES,
+                false /* isBlocked */,
+                mVcnContext,
+                underlyingNetworkTemplates,
+                SUB_GROUP,
+                mSubscriptionSnapshot,
+                null /* currentlySelected */,
+                null /* carrierConfig */);
     }
 
     @Test
@@ -490,24 +501,9 @@ public class NetworkPriorityClassifierTest {
                         mSubscriptionSnapshot));
     }
 
-    private void verifyCalculatePriorityClass(
-            UnderlyingNetworkRecord networkRecord, int expectedIndex) {
-        final int priorityIndex =
-                calculatePriorityClass(
-                        mVcnContext,
-                        networkRecord,
-                        VcnGatewayConnectionConfig.DEFAULT_UNDERLYING_NETWORK_TEMPLATES,
-                        SUB_GROUP,
-                        mSubscriptionSnapshot,
-                        null /* currentlySelected */,
-                        null /* carrierConfig */);
-
-        assertEquals(expectedIndex, priorityIndex);
-    }
-
     @Test
     public void testCalculatePriorityClass() throws Exception {
-        verifyCalculatePriorityClass(mCellNetworkRecord, 2);
+        assertEquals(2, mCellNetworkRecord.priorityClass);
     }
 
     @Test
@@ -518,9 +514,11 @@ public class NetworkPriorityClassifierTest {
                         .setSignalStrength(WIFI_RSSI_LOW)
                         .setSsid(SSID)
                         .build();
+
         final UnderlyingNetworkRecord wifiNetworkRecord =
-                new UnderlyingNetworkRecord(mNetwork, nc, LINK_PROPERTIES, false /* isBlocked */);
+                getTestNetworkRecord(
+                        nc, VcnGatewayConnectionConfig.DEFAULT_UNDERLYING_NETWORK_TEMPLATES);
 
-        verifyCalculatePriorityClass(wifiNetworkRecord, PRIORITY_ANY);
+        assertEquals(PRIORITY_ANY, wifiNetworkRecord.priorityClass);
     }
 }
diff --git a/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java b/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
index fad9669911bb..73d0a280916a 100644
--- a/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
+++ b/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
@@ -21,7 +21,9 @@ import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.WI
 import static com.android.server.vcn.routeselection.NetworkPriorityClassifier.WIFI_EXIT_RSSI_THRESHOLD_DEFAULT;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.any;
@@ -321,16 +323,30 @@ public class UnderlyingNetworkControllerTest {
                 .unregisterNetworkCallback(any(UnderlyingNetworkListener.class));
     }
 
+    private static UnderlyingNetworkRecord getTestNetworkRecord(
+            Network network,
+            NetworkCapabilities networkCapabilities,
+            LinkProperties linkProperties,
+            boolean isBlocked) {
+        return new UnderlyingNetworkRecord(
+                network,
+                networkCapabilities,
+                linkProperties,
+                isBlocked,
+                false /* isSelected */,
+                0 /* priorityClass */);
+    }
+
     @Test
     public void testUnderlyingNetworkRecordEquals() {
         UnderlyingNetworkRecord recordA =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         INITIAL_NETWORK_CAPABILITIES,
                         INITIAL_LINK_PROPERTIES,
                         false /* isBlocked */);
         UnderlyingNetworkRecord recordB =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         INITIAL_NETWORK_CAPABILITIES,
                         INITIAL_LINK_PROPERTIES,
@@ -338,12 +354,24 @@ public class UnderlyingNetworkControllerTest {
         UnderlyingNetworkRecord recordC =
                 new UnderlyingNetworkRecord(
                         mNetwork,
+                        INITIAL_NETWORK_CAPABILITIES,
+                        INITIAL_LINK_PROPERTIES,
+                        false /* isBlocked */,
+                        true /* isSelected */,
+                        -1 /* priorityClass */);
+        UnderlyingNetworkRecord recordD =
+                getTestNetworkRecord(
+                        mNetwork,
                         UPDATED_NETWORK_CAPABILITIES,
                         UPDATED_LINK_PROPERTIES,
                         false /* isBlocked */);
 
         assertEquals(recordA, recordB);
-        assertNotEquals(recordA, recordC);
+        assertEquals(recordA, recordC);
+        assertNotEquals(recordA, recordD);
+
+        assertTrue(UnderlyingNetworkRecord.isEqualIncludingPriorities(recordA, recordB));
+        assertFalse(UnderlyingNetworkRecord.isEqualIncludingPriorities(recordA, recordC));
     }
 
     @Test
@@ -366,6 +394,10 @@ public class UnderlyingNetworkControllerTest {
                 .build();
     }
 
+    private void verifyOnSelectedUnderlyingNetworkChanged(UnderlyingNetworkRecord expectedRecord) {
+        verify(mNetworkControllerCb).onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+    }
+
     private UnderlyingNetworkListener verifyRegistrationOnAvailableAndGetCallback(
             NetworkCapabilities networkCapabilities) {
         verify(mConnectivityManager)
@@ -384,12 +416,12 @@ public class UnderlyingNetworkControllerTest {
         cb.onBlockedStatusChanged(mNetwork, false /* isFalse */);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         responseNetworkCaps,
                         INITIAL_LINK_PROPERTIES,
                         false /* isBlocked */);
-        verify(mNetworkControllerCb).onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
         return cb;
     }
 
@@ -402,12 +434,12 @@ public class UnderlyingNetworkControllerTest {
         cb.onCapabilitiesChanged(mNetwork, responseNetworkCaps);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         responseNetworkCaps,
                         INITIAL_LINK_PROPERTIES,
                         false /* isBlocked */);
-        verify(mNetworkControllerCb).onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
     }
 
     @Test
@@ -417,12 +449,12 @@ public class UnderlyingNetworkControllerTest {
         cb.onLinkPropertiesChanged(mNetwork, UPDATED_LINK_PROPERTIES);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         buildResponseNwCaps(INITIAL_NETWORK_CAPABILITIES, INITIAL_SUB_IDS),
                         UPDATED_LINK_PROPERTIES,
                         false /* isBlocked */);
-        verify(mNetworkControllerCb).onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
     }
 
     @Test
@@ -434,18 +466,16 @@ public class UnderlyingNetworkControllerTest {
         cb.onCapabilitiesChanged(mNetwork, responseNetworkCaps);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         responseNetworkCaps,
                         INITIAL_LINK_PROPERTIES,
                         false /* isBlocked */);
-        verify(mNetworkControllerCb, times(1))
-                .onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
         // onSelectedUnderlyingNetworkChanged() won't be fired twice if network capabilities doesn't
         // change.
         cb.onCapabilitiesChanged(mNetwork, responseNetworkCaps);
-        verify(mNetworkControllerCb, times(1))
-                .onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
     }
 
     @Test
@@ -458,18 +488,16 @@ public class UnderlyingNetworkControllerTest {
         cb.onCapabilitiesChanged(mNetwork, responseNetworkCaps);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         responseNetworkCaps,
                         INITIAL_LINK_PROPERTIES,
                         false /* isBlocked */);
-        verify(mNetworkControllerCb, times(1))
-                .onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
         // onSelectedUnderlyingNetworkChanged() won't be fired twice if network capabilities doesn't
         // change.
         cb.onCapabilitiesChanged(mNetwork, responseNetworkCaps);
-        verify(mNetworkControllerCb, times(1))
-                .onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
     }
 
     @Test
@@ -479,12 +507,12 @@ public class UnderlyingNetworkControllerTest {
         cb.onBlockedStatusChanged(mNetwork, true /* isBlocked */);
 
         UnderlyingNetworkRecord expectedRecord =
-                new UnderlyingNetworkRecord(
+                getTestNetworkRecord(
                         mNetwork,
                         buildResponseNwCaps(INITIAL_NETWORK_CAPABILITIES, INITIAL_SUB_IDS),
                         INITIAL_LINK_PROPERTIES,
                         true /* isBlocked */);
-        verify(mNetworkControllerCb).onSelectedUnderlyingNetworkChanged(eq(expectedRecord));
+        verifyOnSelectedUnderlyingNetworkChanged(expectedRecord);
     }
 
     @Test