From 82b6839c02e7f3bc63bd8a6710db4e6dd7988ef2 Mon Sep 17 00:00:00 2001
From: Chiachang Wang <chiachangwang@google.com>
Date: Mon, 19 Feb 2024 09:31:10 +0000
Subject: [PATCH] Add methods to support verifying captive portal launch

Test: atest FrameworksNetTests
Bug: 325984919
Change-Id: If76a216eca8527b4a7a7f7776d09f66e62daea06
---
 .../base/CSAgentWrapper.kt                    | 37 +++++++++++++++++--
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
index d41c7428d6..d7343b1e3b 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
@@ -19,6 +19,8 @@ package com.android.server
 import android.content.Context
 import android.net.ConnectivityManager
 import android.net.INetworkMonitor
+import android.net.INetworkMonitor.NETWORK_VALIDATION_PROBE_DNS
+import android.net.INetworkMonitor.NETWORK_VALIDATION_PROBE_HTTP
 import android.net.INetworkMonitorCallbacks
 import android.net.LinkProperties
 import android.net.LocalNetworkConfig
@@ -75,10 +77,15 @@ class CSAgentWrapper(
 ) : TestableNetworkCallback.HasNetwork {
     private val TAG = "CSAgent${nextAgentId()}"
     private val VALIDATION_RESULT_INVALID = 0
+    private val NO_PROBE_RESULT = 0
     private val VALIDATION_TIMESTAMP = 1234L
     private val agent: NetworkAgent
     private val nmCallbacks: INetworkMonitorCallbacks
     val networkMonitor = mock<INetworkMonitor>()
+    private var nmValidationRedirectUrl: String? = null
+    private var nmValidationResult = NO_PROBE_RESULT
+    private var nmProbesCompleted = NO_PROBE_RESULT
+    private var nmProbesSucceeded = NO_PROBE_RESULT
 
     override val network: Network get() = agent.network!!
 
@@ -120,10 +127,10 @@ class CSAgentWrapper(
         }
         nmCallbacks.notifyProbeStatusChanged(0 /* completed */, 0 /* succeeded */)
         val p = NetworkTestResultParcelable()
-        p.result = VALIDATION_RESULT_INVALID
-        p.probesAttempted = 0
-        p.probesSucceeded = 0
-        p.redirectUrl = null
+        p.result = nmValidationResult
+        p.probesAttempted = nmProbesCompleted
+        p.probesSucceeded = nmProbesSucceeded
+        p.redirectUrl = nmValidationRedirectUrl
         p.timestampMillis = VALIDATION_TIMESTAMP
         nmCallbacks.notifyNetworkTestedWithExtras(p)
     }
@@ -171,4 +178,26 @@ class CSAgentWrapper(
 
     fun sendLocalNetworkConfig(lnc: LocalNetworkConfig) = agent.sendLocalNetworkConfig(lnc)
     fun sendNetworkCapabilities(nc: NetworkCapabilities) = agent.sendNetworkCapabilities(nc)
+
+    fun connectWithCaptivePortal(redirectUrl: String) {
+        setCaptivePortal(redirectUrl)
+        connect()
+    }
+
+    fun setProbesStatus(probesCompleted: Int, probesSucceeded: Int) {
+        nmProbesCompleted = probesCompleted
+        nmProbesSucceeded = probesSucceeded
+    }
+
+    fun setCaptivePortal(redirectUrl: String) {
+        nmValidationResult = VALIDATION_RESULT_INVALID
+        nmValidationRedirectUrl = redirectUrl
+        // Suppose the portal is found when NetworkMonitor probes NETWORK_VALIDATION_PROBE_HTTP
+        // in the beginning. Because NETWORK_VALIDATION_PROBE_HTTP is the decisive probe for captive
+        // portal, considering the NETWORK_VALIDATION_PROBE_HTTPS hasn't probed yet and set only
+        // DNS and HTTP probes completed.
+        setProbesStatus(
+            NETWORK_VALIDATION_PROBE_DNS or NETWORK_VALIDATION_PROBE_HTTP /* probesCompleted */,
+            VALIDATION_RESULT_INVALID /* probesSucceeded */)
+    }
 }