From 16b160bfcb020d88c2fbfb50f36486d78150ef5f Mon Sep 17 00:00:00 2001
From: Junyu Lai <junyulai@google.com>
Date: Tue, 22 Jun 2021 01:21:24 +0000
Subject: [PATCH 1/2] Add OnSetWarningAndLimit into
 TestableNetworkStatsProvider

This change also address left comments at ag/15010113.

Test: m gts && atest \
      GtsNetworkStackHostTestCases:NetworkStatsHostTest#testNetworkStatsProvider
Bug: 191327585
Merged-In: I100668cc6d4cecbaff29d027d37d313dd48854f3
Change-Id: Ib1f3820fe8502b939d36cef29bdb8dea428604b5
  (cherry-picked from ag/15061296)
---
 .../testutils/TestableNetworkStatsProvider.kt       | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
index d034a7d509..8477151b7e 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
@@ -30,7 +30,12 @@ open class TestableNetworkStatsProvider(
 ) : NetworkStatsProvider() {
     sealed class CallbackType {
         data class OnRequestStatsUpdate(val token: Int) : CallbackType()
-        data class OnSetLimit(val iface: String?, val quotaBytes: Long) : CallbackType()
+        data class OnSetWarningAndLimit(
+            val iface: String?,
+            val warningBytes: Long,
+            val limitBytes: Long
+        ) : CallbackType()
+        data class OnSetLimit(val iface: String?, val limitBytes: Long) : CallbackType()
         data class OnSetAlert(val quotaBytes: Long) : CallbackType()
     }
 
@@ -42,6 +47,10 @@ open class TestableNetworkStatsProvider(
         history.add(CallbackType.OnRequestStatsUpdate(token))
     }
 
+    override fun onSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        history.add(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes))
+    }
+
     override fun onSetLimit(iface: String, quotaBytes: Long) {
         history.add(CallbackType.OnSetLimit(iface, quotaBytes))
     }
@@ -56,7 +65,7 @@ open class TestableNetworkStatsProvider(
         if (token != TOKEN_ANY) {
             assertEquals(token, event.token)
         }
-        return token
+        return event.token
     }
 
     fun expectOnSetLimit(iface: String?, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {

From 749459c3bd81d2d516067ad44e11295e9b932371 Mon Sep 17 00:00:00 2001
From: Junyu Lai <junyulai@google.com>
Date: Fri, 1 Oct 2021 09:52:36 +0000
Subject: [PATCH 2/2] Add getter for quotaBytes

Since the counter part test is in internal branch only.
Change the member of OnSetLimit makes it fail to build on the
rvc+aosp branch.
Thus, add a getter for backward compatibility since the rvc test
does not recognize limitBytes.

This change also addresses some nullability confusions that are
not going to happen.

Test: m gts && atest \
          GtsNetworkStackHostTestCases:NetworkStatsHostTest#testNetworkStatsProvider
      on R & S codebase.
Bug: 191327585
Change-Id: I91dc8cb2297ca503c49b78bf1102cc5572d16530
---
 .../android/testutils/TestableNetworkStatsProvider.kt  | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
index 8477151b7e..be5c9b2eab 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
@@ -31,11 +31,15 @@ open class TestableNetworkStatsProvider(
     sealed class CallbackType {
         data class OnRequestStatsUpdate(val token: Int) : CallbackType()
         data class OnSetWarningAndLimit(
-            val iface: String?,
+            val iface: String,
             val warningBytes: Long,
             val limitBytes: Long
         ) : CallbackType()
-        data class OnSetLimit(val iface: String?, val limitBytes: Long) : CallbackType()
+        data class OnSetLimit(val iface: String, val limitBytes: Long) : CallbackType() {
+            // Add getter for backward compatibility since old tests do not recognize limitBytes.
+            val quotaBytes: Long
+                get() = limitBytes
+        }
         data class OnSetAlert(val quotaBytes: Long) : CallbackType()
     }
 
@@ -68,7 +72,7 @@ open class TestableNetworkStatsProvider(
         return event.token
     }
 
-    fun expectOnSetLimit(iface: String?, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {
+    fun expectOnSetLimit(iface: String, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {
         assertEquals(CallbackType.OnSetLimit(iface, quotaBytes), history.poll(timeout))
     }