From 00a691d1a5a19288d7be30528615d633d7ae0014 Mon Sep 17 00:00:00 2001
From: David Su <dysu@google.com>
Date: Fri, 28 Feb 2020 15:20:15 -0800
Subject: [PATCH] CTS: Add test for WifiManager#getCountryCode()

Bug: 150236894
Test: atest android.net.wifi.cts.WifiManagerTest
Change-Id: I4c77dc63a3c711aa86e3465dbf5dc66196dfc864
---
 .../android/net/wifi/cts/WifiManagerTest.java | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
index c836af5eae..4a8edeed1a 100644
--- a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
@@ -50,6 +50,7 @@ import android.os.UserHandle;
 import android.platform.test.annotations.AppModeFull;
 import android.provider.Settings;
 import android.support.test.uiautomator.UiDevice;
+import android.telephony.TelephonyManager;
 import android.test.AndroidTestCase;
 import android.text.TextUtils;
 import android.util.ArraySet;
@@ -64,6 +65,7 @@ import com.android.compatibility.common.util.SystemUtil;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.util.List;
+import java.util.Locale;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executor;
@@ -1543,4 +1545,29 @@ public class WifiManagerTest extends AndroidTestCase {
             uiAutomation.dropShellPermissionIdentity();
         }
     }
+
+    /**
+     * Test that the wifi country code is either null, or a length-2 string.
+     */
+    public void testGetCountryCode() throws Exception {
+        if (!WifiFeature.isWifiSupported(getContext())) {
+            // skip the test if WiFi is not supported
+            return;
+        }
+
+        String wifiCountryCode = ShellIdentityUtils.invokeWithShellPermissions(
+                mWifiManager::getCountryCode);
+
+        if (wifiCountryCode == null) {
+            return;
+        }
+        assertEquals(2, wifiCountryCode.length());
+
+        // assert that the country code is all uppercase
+        assertEquals(wifiCountryCode.toUpperCase(Locale.US), wifiCountryCode);
+
+        String telephonyCountryCode = getContext().getSystemService(TelephonyManager.class)
+                .getNetworkCountryIso();
+        assertEquals(telephonyCountryCode, wifiCountryCode.toLowerCase(Locale.US));
+    }
 }