diff --git a/tests/unit/java/android/net/BpfNetMapsReaderTest.kt b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
index 86a0acb..258e422 100644
--- a/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
+++ b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
@@ -17,6 +17,8 @@
 package android.net
 
 import android.net.BpfNetMapsConstants.DOZABLE_MATCH
+import android.net.BpfNetMapsConstants.HAPPY_BOX_MATCH
+import android.net.BpfNetMapsConstants.PENALTY_BOX_MATCH
 import android.net.BpfNetMapsConstants.STANDBY_MATCH
 import android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY
 import android.net.BpfNetMapsUtils.getMatchByFirewallChain
@@ -36,6 +38,7 @@
 
 private const val TEST_UID1 = 1234
 private const val TEST_UID2 = TEST_UID1 + 1
+private const val TEST_UID3 = TEST_UID2 + 1
 private const val NO_IIF = 0
 
 // pre-T devices does not support Bpf.
@@ -101,23 +104,26 @@
         testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(newConfig))
     }
 
+    fun isUidNetworkingBlocked(uid: Int, metered: Boolean = false, dataSaver: Boolean = false) =
+            bpfNetMapsReader.isUidNetworkingBlocked(uid, metered, dataSaver)
+
     @Test
     fun testIsUidNetworkingBlockedByFirewallChains_allowChain() {
         // With everything disabled by default, verify the return value is false.
         testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
-        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertFalse(isUidNetworkingBlocked(TEST_UID1))
 
         // Enable dozable chain but does not provide allowed list. Verify the network is blocked
         // for all uids.
         mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_DOZABLE, true)
-        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
-        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+        assertTrue(isUidNetworkingBlocked(TEST_UID1))
+        assertTrue(isUidNetworkingBlocked(TEST_UID2))
 
         // Add uid1 to dozable allowed list. Verify the network is not blocked for uid1, while
         // uid2 is blocked.
         testUidOwnerMap.updateEntry(S32(TEST_UID1), UidOwnerValue(NO_IIF, DOZABLE_MATCH))
-        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
-        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+        assertFalse(isUidNetworkingBlocked(TEST_UID1))
+        assertTrue(isUidNetworkingBlocked(TEST_UID2))
     }
 
     @Test
@@ -126,14 +132,14 @@
         // for all uids.
         testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
         mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_STANDBY, true)
-        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
-        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+        assertFalse(isUidNetworkingBlocked(TEST_UID1))
+        assertFalse(isUidNetworkingBlocked(TEST_UID2))
 
         // Add uid1 to standby allowed list. Verify the network is blocked for uid1, while
         // uid2 is not blocked.
         testUidOwnerMap.updateEntry(S32(TEST_UID1), UidOwnerValue(NO_IIF, STANDBY_MATCH))
-        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
-        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+        assertTrue(isUidNetworkingBlocked(TEST_UID1))
+        assertFalse(isUidNetworkingBlocked(TEST_UID2))
     }
 
     @Test
@@ -143,6 +149,54 @@
         testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
         mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_POWERSAVE, true)
         mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_STANDBY, true)
-        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertTrue(isUidNetworkingBlocked(TEST_UID1))
+    }
+
+    @IgnoreUpTo(VERSION_CODES.S_V2)
+    @Test
+    fun testIsUidNetworkingBlockedByDataSaver() {
+        // With everything disabled by default, verify the return value is false.
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
+        assertFalse(isUidNetworkingBlocked(TEST_UID1, metered = true))
+
+        // Add uid1 to penalty box, verify the network is blocked for uid1, while uid2 is not
+        // affected.
+        testUidOwnerMap.updateEntry(S32(TEST_UID1), UidOwnerValue(NO_IIF, PENALTY_BOX_MATCH))
+        assertTrue(isUidNetworkingBlocked(TEST_UID1, metered = true))
+        assertFalse(isUidNetworkingBlocked(TEST_UID2, metered = true))
+
+        // Enable data saver, verify the network is blocked for uid1, uid2, but uid3 in happy box
+        // is not affected.
+        testUidOwnerMap.updateEntry(S32(TEST_UID3), UidOwnerValue(NO_IIF, HAPPY_BOX_MATCH))
+        assertTrue(isUidNetworkingBlocked(TEST_UID1, metered = true, dataSaver = true))
+        assertTrue(isUidNetworkingBlocked(TEST_UID2, metered = true, dataSaver = true))
+        assertFalse(isUidNetworkingBlocked(TEST_UID3, metered = true, dataSaver = true))
+
+        // Add uid1 to happy box as well, verify nothing is changed because penalty box has higher
+        // priority.
+        testUidOwnerMap.updateEntry(
+            S32(TEST_UID1),
+            UidOwnerValue(NO_IIF, PENALTY_BOX_MATCH or HAPPY_BOX_MATCH)
+        )
+        assertTrue(isUidNetworkingBlocked(TEST_UID1, metered = true, dataSaver = true))
+        assertTrue(isUidNetworkingBlocked(TEST_UID2, metered = true, dataSaver = true))
+        assertFalse(isUidNetworkingBlocked(TEST_UID3, metered = true, dataSaver = true))
+
+        // Enable doze mode, verify uid3 is blocked even if it is in happy box.
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_DOZABLE, true)
+        assertTrue(isUidNetworkingBlocked(TEST_UID1, metered = true, dataSaver = true))
+        assertTrue(isUidNetworkingBlocked(TEST_UID2, metered = true, dataSaver = true))
+        assertTrue(isUidNetworkingBlocked(TEST_UID3, metered = true, dataSaver = true))
+
+        // Disable doze mode and data saver, only uid1 which is in penalty box is blocked.
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_DOZABLE, false)
+        assertTrue(isUidNetworkingBlocked(TEST_UID1, metered = true))
+        assertFalse(isUidNetworkingBlocked(TEST_UID2, metered = true))
+        assertFalse(isUidNetworkingBlocked(TEST_UID3, metered = true))
+
+        // Make the network non-metered, nothing is blocked.
+        assertFalse(isUidNetworkingBlocked(TEST_UID1))
+        assertFalse(isUidNetworkingBlocked(TEST_UID2))
+        assertFalse(isUidNetworkingBlocked(TEST_UID3))
     }
 }
diff --git a/tests/unit/java/android/net/ConnectivityManagerTest.java b/tests/unit/java/android/net/ConnectivityManagerTest.java
index 45a9dbc..b8c5447 100644
--- a/tests/unit/java/android/net/ConnectivityManagerTest.java
+++ b/tests/unit/java/android/net/ConnectivityManagerTest.java
@@ -16,6 +16,13 @@
 
 package android.net;
 
+import static android.content.Context.RECEIVER_NOT_EXPORTED;
+import static android.content.pm.ApplicationInfo.FLAG_PERSISTENT;
+import static android.content.pm.ApplicationInfo.FLAG_SYSTEM;
+import static android.net.ConnectivityManager.ACTION_RESTRICT_BACKGROUND_CHANGED;
+import static android.net.ConnectivityManager.RESTRICT_BACKGROUND_STATUS_DISABLED;
+import static android.net.ConnectivityManager.RESTRICT_BACKGROUND_STATUS_ENABLED;
+import static android.net.ConnectivityManager.RESTRICT_BACKGROUND_STATUS_WHITELISTED;
 import static android.net.ConnectivityManager.TYPE_NONE;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_CBS;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_DUN;
@@ -39,6 +46,7 @@
 
 import static com.android.testutils.MiscAsserts.assertThrows;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
@@ -51,6 +59,7 @@
 import static org.mockito.Mockito.after;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -61,7 +70,10 @@
 
 import android.app.PendingIntent;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ApplicationInfo;
+import android.net.ConnectivityManager.DataSaverStatusTracker;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.os.Build.VERSION_CODES;
 import android.os.Bundle;
@@ -95,6 +107,7 @@
 
     @Mock Context mCtx;
     @Mock IConnectivityManager mService;
+    @Mock NetworkPolicyManager mNpm;
 
     @Before
     public void setUp() {
@@ -510,4 +523,54 @@
         assertNull("ConnectivityManager weak reference still not null after " + attempts
                     + " attempts", ref.get());
     }
+
+    @Test
+    public void testDataSaverStatusTracker() {
+        mockService(NetworkPolicyManager.class, Context.NETWORK_POLICY_SERVICE, mNpm);
+        // Mock proper application info.
+        doReturn(mCtx).when(mCtx).getApplicationContext();
+        final ApplicationInfo mockAppInfo = new ApplicationInfo();
+        mockAppInfo.flags = FLAG_PERSISTENT | FLAG_SYSTEM;
+        doReturn(mockAppInfo).when(mCtx).getApplicationInfo();
+        // Enable data saver.
+        doReturn(RESTRICT_BACKGROUND_STATUS_ENABLED).when(mNpm)
+                .getRestrictBackgroundStatus(anyInt());
+
+        final DataSaverStatusTracker tracker = new DataSaverStatusTracker(mCtx);
+        // Verify the data saver status is correct right after initialization.
+        assertTrue(tracker.getDataSaverEnabled());
+
+        // Verify the tracker register receiver with expected intent filter.
+        final ArgumentCaptor<IntentFilter> intentFilterCaptor =
+                ArgumentCaptor.forClass(IntentFilter.class);
+        verify(mCtx).registerReceiver(
+                any(), intentFilterCaptor.capture(), eq(RECEIVER_NOT_EXPORTED));
+        assertEquals(ACTION_RESTRICT_BACKGROUND_CHANGED,
+                intentFilterCaptor.getValue().getAction(0));
+
+        // Mock data saver status changed event and verify the tracker tracks the
+        // status accordingly.
+        doReturn(RESTRICT_BACKGROUND_STATUS_DISABLED).when(mNpm)
+                .getRestrictBackgroundStatus(anyInt());
+        tracker.onReceive(mCtx, new Intent(ACTION_RESTRICT_BACKGROUND_CHANGED));
+        assertFalse(tracker.getDataSaverEnabled());
+
+        doReturn(RESTRICT_BACKGROUND_STATUS_WHITELISTED).when(mNpm)
+                .getRestrictBackgroundStatus(anyInt());
+        tracker.onReceive(mCtx, new Intent(ACTION_RESTRICT_BACKGROUND_CHANGED));
+        assertTrue(tracker.getDataSaverEnabled());
+    }
+
+    private <T> void mockService(Class<T> clazz, String name, T service) {
+        doReturn(service).when(mCtx).getSystemService(name);
+        doReturn(name).when(mCtx).getSystemServiceName(clazz);
+
+        // If the test suite uses the inline mock maker library, such as for coverage tests,
+        // then the final version of getSystemService must also be mocked, as the real
+        // method will not be called by the test and null object is returned since no mock.
+        // Otherwise, mocking a final method will fail the test.
+        if (mCtx.getSystemService(clazz) == null) {
+            doReturn(service).when(mCtx).getSystemService(clazz);
+        }
+    }
 }