diff --git a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt
index 84922685faf002b91a38ffa4dc5c6a3cb625347e..e80548b3d20e698e6b6b865a93ecac5a0768d76b 100644
--- a/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt
+++ b/tests/benchmark/src/android/net/netstats/benchmarktests/NetworkStatsTest.kt
@@ -28,7 +28,6 @@ import com.android.server.net.NetworkStatsRecorder
 import java.io.BufferedInputStream
 import java.io.DataInputStream
 import java.io.File
-import java.io.FileInputStream
 import java.io.FileOutputStream
 import java.nio.file.Files
 import java.util.concurrent.TimeUnit
@@ -44,6 +43,8 @@ import org.mockito.Mockito.mock
 class NetworkStatsTest {
     companion object {
         private val DEFAULT_BUFFER_SIZE = 8192
+        private val FILE_CACHE_WARM_UP_REPEAT_COUNT = 10
+        private val TEST_REPEAT_COUNT = 10
         private val UID_COLLECTION_BUCKET_DURATION_MS = TimeUnit.HOURS.toMillis(2)
         private val UID_RECORDER_ROTATE_AGE_MS = TimeUnit.DAYS.toMillis(15)
         private val UID_RECORDER_DELETE_AGE_MS = TimeUnit.DAYS.toMillis(90)
@@ -63,14 +64,14 @@ class NetworkStatsTest {
         }
 
         // Test results shows the test cases who read the file first will take longer time to
-        // execute, and reading time getting shorter each time. Read files several times prior to
-        // tests to minimize the impact. This cannot live in setUp() since the time
-        // spent on the file reading will be attributed to the time spent on the individual
-        // test case.
+        // execute, and reading time getting shorter each time due to file caching mechanism.
+        // Read files several times prior to tests to minimize the impact.
+        // This cannot live in setUp() since the time spent on the file reading will be
+        // attributed to the time spent on the individual test case.
         @JvmStatic
         @BeforeClass
         fun setUpOnce() {
-            for (i in 1..10) {
+            repeat(FILE_CACHE_WARM_UP_REPEAT_COUNT) {
                 val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
                 for (file in uidTestFiles) {
                     readFile(file, collection)
@@ -78,20 +79,19 @@ class NetworkStatsTest {
             }
         }
 
-        private fun getInputStreamForResource(resourceId: Int): DataInputStream {
-            return DataInputStream(
+        private fun getInputStreamForResource(resourceId: Int): DataInputStream =
+            DataInputStream(
                 InstrumentationRegistry.getContext()
                     .getResources().openRawResource(resourceId)
             )
-        }
 
         private fun unzipToTempDir(zis: ZipInputStream): File {
             val statsDir =
                 Files.createTempDirectory(NetworkStatsTest::class.simpleName).toFile()
-            while (true) {
-                val entryName = zis.nextEntry?.name ?: break
-                val file = File(statsDir, entryName)
-                FileOutputStream(file).use { zis.copyTo(it, DEFAULT_BUFFER_SIZE) }
+            generateSequence { zis.nextEntry }.forEach { entry ->
+                FileOutputStream(File(statsDir, entry.name)).use {
+                    zis.copyTo(it, DEFAULT_BUFFER_SIZE)
+                }
             }
             return statsDir
         }
@@ -99,20 +99,23 @@ class NetworkStatsTest {
         // List [xt|uid|uid_tag].<start>-<end> files under the given directory.
         private fun getSortedListForPrefix(statsDir: File, prefix: String): List<File> {
             assertTrue(statsDir.exists())
-            return (statsDir.list() ?: arrayOf()).mapNotNull {
-                if (it.startsWith("$prefix.")) File(statsDir, it) else null
-            }.sorted()
+            return statsDir.list() { dir, name -> name.startsWith("$prefix.") }
+                .orEmpty()
+                .map { it -> File(statsDir, it) }
+                .sorted()
         }
 
         private fun readFile(file: File, reader: Reader) =
-            BufferedInputStream(FileInputStream(file)).use {
+            BufferedInputStream(file.inputStream()).use {
                 reader.read(it)
             }
     }
 
     @Test
     fun testReadCollection_manyUids() {
-        for (i in 1..10) {
+        // The file cache is warmed up by the @BeforeClass method, so now the test can repeat
+        // this a number of time to have a stable number.
+        repeat(TEST_REPEAT_COUNT) {
             val collection = NetworkStatsCollection(UID_COLLECTION_BUCKET_DURATION_MS)
             for (file in uidTestFiles) {
                 readFile(file, collection)
@@ -122,13 +125,15 @@ class NetworkStatsTest {
 
     @Test
     fun testReadFromRecorder_manyUids() {
-        for (i in 1..10) {
+        val mockObserver = mock<NonMonotonicObserver<String>>()
+        val mockDropBox = mock<DropBoxManager>()
+        repeat(TEST_REPEAT_COUNT) {
             val recorder = NetworkStatsRecorder(
                 FileRotator(
                     testFilesDir, PREFIX_UID, UID_RECORDER_ROTATE_AGE_MS, UID_RECORDER_DELETE_AGE_MS
                 ),
-                mock<NonMonotonicObserver<String>>(),
-                mock(DropBoxManager::class.java),
+                mockObserver,
+                mockDropBox,
                 PREFIX_UID,
                 UID_COLLECTION_BUCKET_DURATION_MS,
                 false /* includeTags */,
@@ -138,5 +143,5 @@ class NetworkStatsTest {
         }
     }
 
-    inline fun <reified T : Any> mock(): T = mock(T::class.java)
+    inline fun <reified T> mock(): T = mock(T::class.java)
 }