From be9447be582f289058b313463a89527f6df7f9ed Mon Sep 17 00:00:00 2001
From: Chalard Jean <jchalard@google.com>
Date: Tue, 7 Apr 2020 19:15:37 +0900
Subject: [PATCH] Add instrumentation to test NetworkAgent.

Test: NetworkAgentTest, new tests using this API
Bug: 139268426
Change-Id: Ia83b1c896df63bb18e2aa4b74d6cc09eba990eb5
---
 core/java/android/net/NetworkAgent.java | 33 +++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/core/java/android/net/NetworkAgent.java b/core/java/android/net/NetworkAgent.java
index 31d7d08209..65e772cb5e 100644
--- a/core/java/android/net/NetworkAgent.java
+++ b/core/java/android/net/NetworkAgent.java
@@ -25,12 +25,14 @@ import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.os.Build;
 import android.os.Bundle;
+import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
 import android.os.Messenger;
 import android.util.Log;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.AsyncChannel;
 import com.android.internal.util.Protocol;
 
@@ -571,6 +573,37 @@ public abstract class NetworkAgent {
         return mNetwork;
     }
 
+    /**
+     * Register this network agent with a testing harness.
+     *
+     * The returned Messenger sends messages to the Handler. This allows a test to send
+     * this object {@code CMD_*} messages as if they came from ConnectivityService, which
+     * is useful for testing the behavior.
+     *
+     * @hide
+     */
+    public Messenger registerForTest(final Network network) {
+        log("Registering NetworkAgent for test");
+        synchronized (mRegisterLock) {
+            mNetwork = network;
+            mInitialConfiguration = null;
+        }
+        return new Messenger(mHandler);
+    }
+
+    /**
+     * Waits for the handler to be idle.
+     * This is useful for testing, and has smaller scope than an accessor to mHandler.
+     * TODO : move the implementation in common library with the tests
+     * @hide
+     */
+    @VisibleForTesting
+    public boolean waitForIdle(final long timeoutMs) {
+        final ConditionVariable cv = new ConditionVariable(false);
+        mHandler.post(cv::open);
+        return cv.block(timeoutMs);
+    }
+
     /**
      * @return The Network associated with this agent, or null if it's not registered yet.
      */