From 8fd6347d356b2e4109a88daec8f746a9ad1b832b Mon Sep 17 00:00:00 2001
From: Chalard Jean <jchalard@google.com>
Date: Thu, 13 Mar 2025 17:13:34 +0900
Subject: [PATCH] Add a method to help mocking registerNetworkAgent

This is a trick to help mocking registerNetworkAgent without
having to import its return type. In the near future the
return type will be changed to a non-public class, which
can't be seen in CTS.

Really this shouldn't be necessary because only Connectivity
tests should mock ConnectivityManager, but the practice is
different and other mockers need this support.

See also aosp/3538922

Test: NetworkAgentTest
Change-Id: I3ac63d95d3830a3b413f9ad12b84ac426b63dd08
---
 .../src/android/net/ConnectivityManager.java     | 16 ++++++++++++++++
 .../net/src/android/net/cts/NetworkAgentTest.kt  | 14 ++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 5d99b747b9..3d7ea6944e 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -1242,6 +1242,22 @@ public class ConnectivityManager {
     @ConnectivityManagerFeature
     private Long mEnabledConnectivityManagerFeatures = null;
 
+    /**
+     * A class to help with mocking ConnectivityManager.
+     * @hide
+     */
+    public static class MockHelpers {
+        /**
+         * Produce an instance of the class returned by
+         * {@link ConnectivityManager#registerNetworkAgent}
+         * @hide
+         */
+        public static Network registerNetworkAgentResult(
+                @Nullable final Network network, @Nullable final INetworkAgentRegistry registry) {
+            return network;
+        }
+    }
+
     private TetheringManager getTetheringManager() {
         synchronized (mTetheringEventCallbacks) {
             if (mTetheringManager == null) {
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 8fcc703c0a..5e035a2723 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -172,6 +172,7 @@ import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.any
+import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.argThat
 import org.mockito.ArgumentMatchers.eq
 import org.mockito.Mockito.doReturn
@@ -1066,7 +1067,20 @@ class NetworkAgentTest {
     fun testAgentStartsInConnecting() {
         val mockContext = mock(Context::class.java)
         val mockCm = mock(ConnectivityManager::class.java)
+        val mockedResult = ConnectivityManager.MockHelpers.registerNetworkAgentResult(
+            mock(Network::class.java),
+            mock(INetworkAgentRegistry::class.java)
+        )
         doReturn(mockCm).`when`(mockContext).getSystemService(Context.CONNECTIVITY_SERVICE)
+        doReturn(mockedResult).`when`(mockCm).registerNetworkAgent(
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            anyInt()
+        )
         val agent = createNetworkAgent(mockContext)
         agent.register()
         verify(mockCm).registerNetworkAgent(