diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 76c0c38..34bcf3f 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -198,7 +198,10 @@
     @Mock
     private LocationPermissionChecker mLocationPermissionChecker;
     private @Mock IBpfMap<U32, U8> mUidCounterSetMap;
-    private @Mock NetworkStatsService.TagStatsDeleter mTagStatsDeleter;
+    private @Mock IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap;
+    private @Mock IBpfMap<StatsMapKey, StatsMapValue> mStatsMapA;
+    private @Mock IBpfMap<StatsMapKey, StatsMapValue> mStatsMapB;
+    private @Mock IBpfMap<UidStatsMapKey, StatsMapValue> mAppUidStatsMap;
 
     private NetworkStatsService mService;
     private INetworkStatsSession mSession;
@@ -361,8 +364,23 @@
             }
 
             @Override
-            public NetworkStatsService.TagStatsDeleter getTagStatsDeleter() {
-                return mTagStatsDeleter;
+            public IBpfMap<CookieTagMapKey, CookieTagMapValue> getCookieTagMap() {
+                return mCookieTagMap;
+            }
+
+            @Override
+            public IBpfMap<StatsMapKey, StatsMapValue> getStatsMapA() {
+                return mStatsMapA;
+            }
+
+            @Override
+            public IBpfMap<StatsMapKey, StatsMapValue> getStatsMapB() {
+                return mStatsMapB;
+            }
+
+            @Override
+            public IBpfMap<UidStatsMapKey, StatsMapValue> getAppUidStatsMap() {
+                return mAppUidStatsMap;
             }
         };
     }
@@ -702,10 +720,8 @@
         final Intent intent = new Intent(ACTION_UID_REMOVED);
         intent.putExtra(EXTRA_UID, UID_BLUE);
         mServiceContext.sendBroadcast(intent);
-        verify(mTagStatsDeleter).deleteTagData(UID_BLUE);
         intent.putExtra(EXTRA_UID, UID_RED);
         mServiceContext.sendBroadcast(intent);
-        verify(mTagStatsDeleter).deleteTagData(UID_RED);
 
         // existing uid and total should remain unchanged; but removed UID
         // should be gone completely.