Add WiFi country code to Thread country code sources

Bug: b/309357909
Test: Run `atest ThreadNetworkUnitTests`.
Change-Id: I087cc0fc2b6c21fd9d3d4607f30c194214aeb37e
diff --git a/thread/tests/unit/Android.bp b/thread/tests/unit/Android.bp
index c7887bc..2523915 100644
--- a/thread/tests/unit/Android.bp
+++ b/thread/tests/unit/Android.bp
@@ -47,6 +47,7 @@
         "android.test.base",
         "android.test.runner",
         "ServiceConnectivityResources",
+        "framework-wifi",
     ],
     jarjar_rules: ":connectivity-jarjar-rules",
     jni_libs: [
diff --git a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
index a0eff6c..d7082fe 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
@@ -32,6 +32,7 @@
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -45,6 +46,8 @@
 import android.location.LocationListener;
 import android.location.LocationManager;
 import android.net.thread.IOperationReceiver;
+import android.net.wifi.WifiManager;
+import android.net.wifi.WifiManager.ActiveCountryCodeChangedCallback;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -78,6 +81,7 @@
     @Mock Location mLocation;
     @Mock Resources mResources;
     @Mock ConnectivityResources mConnectivityResources;
+    @Mock WifiManager mWifiManager;
 
     private ThreadNetworkCountryCode mThreadNetworkCountryCode;
     private boolean mErrorSetCountryCode;
@@ -85,6 +89,7 @@
     @Captor private ArgumentCaptor<LocationListener> mLocationListenerCaptor;
     @Captor private ArgumentCaptor<Geocoder.GeocodeListener> mGeocodeListenerCaptor;
     @Captor private ArgumentCaptor<IOperationReceiver> mOperationReceiverCaptor;
+    @Captor private ArgumentCaptor<ActiveCountryCodeChangedCallback> mWifiCountryCodeReceiverCaptor;
 
     @Before
     public void setUp() throws Exception {
@@ -118,7 +123,8 @@
                         mLocationManager,
                         mThreadNetworkControllerService,
                         mGeocoder,
-                        mConnectivityResources);
+                        mConnectivityResources,
+                        mWifiManager);
     }
 
     private static Address newAddress(String countryCode) {
@@ -162,6 +168,55 @@
     }
 
     @Test
+    public void wifiCountryCode_bothWifiAndLocationAreAvailable_wifiCountryCodeIsUsed() {
+        mThreadNetworkCountryCode.initialize();
+        verify(mLocationManager)
+                .requestLocationUpdates(
+                        anyString(), anyLong(), anyFloat(), mLocationListenerCaptor.capture());
+        mLocationListenerCaptor.getValue().onLocationChanged(mLocation);
+        verify(mGeocoder)
+                .getFromLocation(
+                        anyDouble(), anyDouble(), anyInt(), mGeocodeListenerCaptor.capture());
+        Address mockAddress = mock(Address.class);
+        when(mockAddress.getCountryCode()).thenReturn(TEST_COUNTRY_CODE_US);
+        List<Address> addresses = List.of(mockAddress);
+        mGeocodeListenerCaptor.getValue().onGeocode(addresses);
+
+        verify(mWifiManager)
+                .registerActiveCountryCodeChangedCallback(
+                        any(), mWifiCountryCodeReceiverCaptor.capture());
+        mWifiCountryCodeReceiverCaptor.getValue().onActiveCountryCodeChanged(TEST_COUNTRY_CODE_CN);
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode()).isEqualTo(TEST_COUNTRY_CODE_CN);
+    }
+
+    @Test
+    public void wifiCountryCode_wifiCountryCodeIsActive_wifiCountryCodeIsUsed() {
+        mThreadNetworkCountryCode.initialize();
+
+        verify(mWifiManager)
+                .registerActiveCountryCodeChangedCallback(
+                        any(), mWifiCountryCodeReceiverCaptor.capture());
+        mWifiCountryCodeReceiverCaptor.getValue().onActiveCountryCodeChanged(TEST_COUNTRY_CODE_US);
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode()).isEqualTo(TEST_COUNTRY_CODE_US);
+    }
+
+    @Test
+    public void wifiCountryCode_wifiCountryCodeIsInactive_defaultCountryCodeIsUsed() {
+        mThreadNetworkCountryCode.initialize();
+        verify(mWifiManager)
+                .registerActiveCountryCodeChangedCallback(
+                        any(), mWifiCountryCodeReceiverCaptor.capture());
+        mWifiCountryCodeReceiverCaptor.getValue().onActiveCountryCodeChanged(TEST_COUNTRY_CODE_US);
+
+        mWifiCountryCodeReceiverCaptor.getValue().onCountryCodeInactive();
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode())
+                .isEqualTo(ThreadNetworkCountryCode.DEFAULT_COUNTRY_CODE);
+    }
+
+    @Test
     public void updateCountryCode_noForceUpdateDefaultCountryCode_noCountryCodeIsUpdated() {
         mThreadNetworkCountryCode.initialize();
         clearInvocations(mThreadNetworkControllerService);