From fc9119df29f67d70c887a67e08ea510d48b5a124 Mon Sep 17 00:00:00 2001
From: Remi NGUYEN VAN <reminv@google.com>
Date: Thu, 19 Mar 2020 11:25:50 +0900
Subject: [PATCH] Add a test API for testing NetworkStack.getService

NetworkStack.getService was introduced to avoid relying on
Context#getSystemService to obtain the NetworkStack binder token.
To allow it to be mocked in tests, a method is introduced to allow tests
to specify a mock NetworkStack token for their own process.

Test: atest NetworkStackTests, using the change
Bug: 151243982
Change-Id: I04058a007f2dfe1044cabeb3ac508404873665ad
---
 api/test-current.txt                    |  1 +
 core/java/android/net/NetworkStack.java | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/api/test-current.txt b/api/test-current.txt
index 95e060bd61079..3bc5cc7ff07b2 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -1817,6 +1817,7 @@ package android.net {
 
   public class NetworkStack {
     method @Nullable public static android.os.IBinder getService();
+    method public static void setServiceForTest(@Nullable android.os.IBinder);
     field public static final String PERMISSION_MAINLINE_NETWORK_STACK = "android.permission.MAINLINE_NETWORK_STACK";
   }
 
diff --git a/core/java/android/net/NetworkStack.java b/core/java/android/net/NetworkStack.java
index a6d52d133ce90..86f3dfd412f64 100644
--- a/core/java/android/net/NetworkStack.java
+++ b/core/java/android/net/NetworkStack.java
@@ -45,6 +45,9 @@ public class NetworkStack {
     public static final String PERMISSION_MAINLINE_NETWORK_STACK =
             "android.permission.MAINLINE_NETWORK_STACK";
 
+    @Nullable
+    private static volatile IBinder sMockService;
+
     /**
      * Get an {@link IBinder} representing the NetworkStack stable AIDL Interface, if registered.
      * @hide
@@ -53,9 +56,23 @@ public class NetworkStack {
     @SystemApi
     @TestApi
     public static IBinder getService() {
+        final IBinder mockService = sMockService;
+        if (mockService != null) return mockService;
         return ServiceManager.getService(Context.NETWORK_STACK_SERVICE);
     }
 
+    /**
+     * Set a mock service for testing, to be returned by future calls to {@link #getService()}.
+     *
+     * <p>Passing a {@code null} {@code mockService} resets {@link #getService()} to normal
+     * behavior.
+     * @hide
+     */
+    @TestApi
+    public static void setServiceForTest(@Nullable IBinder mockService) {
+        sMockService = mockService;
+    }
+
     private NetworkStack() {}
 
     /**