Merge "Fix flaky test: testDataMigration_differentFromFallback" into main
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 3f4113a..c46eada 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -63,6 +63,7 @@
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 import static android.text.format.DateUtils.SECOND_IN_MILLIS;
 
+import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
 import static com.android.net.module.util.NetworkCapabilitiesUtils.getDisplayTransport;
 import static com.android.net.module.util.NetworkStatsUtils.LIMIT_GLOBAL_ALERT;
 
@@ -3248,7 +3249,8 @@
      * Default external settings that read from
      * {@link android.provider.Settings.Global}.
      */
-    private static class DefaultNetworkStatsSettings implements NetworkStatsSettings {
+    @VisibleForTesting(visibility = PRIVATE)
+    static class DefaultNetworkStatsSettings implements NetworkStatsSettings {
         DefaultNetworkStatsSettings() {}
 
         @Override
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 6292d45..9453617 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -79,7 +79,6 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doReturn;
@@ -240,7 +239,9 @@
     private @Mock INetd mNetd;
     private @Mock TetheringManager mTetheringManager;
     private @Mock NetworkStatsFactory mStatsFactory;
-    private @Mock NetworkStatsSettings mSettings;
+    @NonNull
+    private final TestNetworkStatsSettings mSettings =
+            new TestNetworkStatsSettings(HOUR_IN_MILLIS, WEEK_IN_MILLIS);
     private @Mock IBinder mUsageCallbackBinder;
     private TestableUsageCallback mUsageCallback;
     private @Mock AlarmManager mAlarmManager;
@@ -533,7 +534,6 @@
         mStatsDir = null;
 
         mNetd = null;
-        mSettings = null;
 
         mSession.close();
         mService = null;
@@ -1765,7 +1765,7 @@
     }
 
     private void setCombineSubtypeEnabled(boolean enable) {
-        doReturn(enable).when(mSettings).getCombineSubtypeEnabled();
+        mSettings.setCombineSubtypeEnabled(enable);
         mHandler.post(() -> mContentObserver.onChange(false, Settings.Global
                     .getUriFor(Settings.Global.NETSTATS_COMBINE_SUBTYPE_ENABLED)));
         waitForIdle();
@@ -2289,21 +2289,80 @@
         mockSettings(HOUR_IN_MILLIS, WEEK_IN_MILLIS);
     }
 
-    private void mockSettings(long bucketDuration, long deleteAge) throws Exception {
-        doReturn(HOUR_IN_MILLIS).when(mSettings).getPollInterval();
-        doReturn(0L).when(mSettings).getPollDelay();
-        doReturn(true).when(mSettings).getSampleEnabled();
-        doReturn(false).when(mSettings).getCombineSubtypeEnabled();
+    private void mockSettings(long bucketDuration, long deleteAge) {
+        mSettings.setConfig(new Config(bucketDuration, deleteAge, deleteAge));
+    }
 
-        final Config config = new Config(bucketDuration, deleteAge, deleteAge);
-        doReturn(config).when(mSettings).getXtConfig();
-        doReturn(config).when(mSettings).getUidConfig();
-        doReturn(config).when(mSettings).getUidTagConfig();
+    // Note that this object will be accessed from test main thread and service handler thread.
+    // Thus, it has to be thread safe in order to prevent from flakiness.
+    private static class TestNetworkStatsSettings
+            extends NetworkStatsService.DefaultNetworkStatsSettings {
 
-        doReturn(MB_IN_BYTES).when(mSettings).getGlobalAlertBytes(anyLong());
-        doReturn(MB_IN_BYTES).when(mSettings).getXtPersistBytes(anyLong());
-        doReturn(MB_IN_BYTES).when(mSettings).getUidPersistBytes(anyLong());
-        doReturn(MB_IN_BYTES).when(mSettings).getUidTagPersistBytes(anyLong());
+        @NonNull
+        private volatile Config mConfig;
+        private final AtomicBoolean mCombineSubtypeEnabled = new AtomicBoolean();
+
+        TestNetworkStatsSettings(long bucketDuration, long deleteAge) {
+            mConfig = new Config(bucketDuration, deleteAge, deleteAge);
+        }
+
+        void setConfig(@NonNull Config config) {
+            mConfig = config;
+        }
+
+        @Override
+        public long getPollDelay() {
+            return 0L;
+        }
+
+        @Override
+        public long getGlobalAlertBytes(long def) {
+            return MB_IN_BYTES;
+        }
+
+        @Override
+        public Config getXtConfig() {
+            return mConfig;
+        }
+
+        @Override
+        public Config getUidConfig() {
+            return mConfig;
+        }
+
+        @Override
+        public Config getUidTagConfig() {
+            return mConfig;
+        }
+
+        @Override
+        public long getXtPersistBytes(long def) {
+            return MB_IN_BYTES;
+        }
+
+        @Override
+        public long getUidPersistBytes(long def) {
+            return MB_IN_BYTES;
+        }
+
+        @Override
+        public long getUidTagPersistBytes(long def) {
+            return MB_IN_BYTES;
+        }
+
+        @Override
+        public boolean getCombineSubtypeEnabled() {
+            return mCombineSubtypeEnabled.get();
+        }
+
+        public void setCombineSubtypeEnabled(boolean enable) {
+            mCombineSubtypeEnabled.set(enable);
+        }
+
+        @Override
+        public boolean getAugmentEnabled() {
+            return false;
+        }
     }
 
     private void assertStatsFilesExist(boolean exist) {