[WifiManager11axTest] WifiManager 11ax tests.

Bug: 173730855
Test: Verified the changes
Change-Id: Ib48bacab4400c1ced782f24b04d0371b7c129f87
diff --git a/acts_tests/tests/google/wifi/WifiManagerTest.py b/acts_tests/tests/google/wifi/WifiManagerTest.py
index dbb447e..719d2ca 100644
--- a/acts_tests/tests/google/wifi/WifiManagerTest.py
+++ b/acts_tests/tests/google/wifi/WifiManagerTest.py
@@ -71,7 +71,7 @@
         opt_param = [
             "open_network", "reference_networks", "iperf_server_address",
             "wpa_networks", "wep_networks", "iperf_server_port",
-            "coex_unsafe_channels", "coex_restrictions"
+            "coex_unsafe_channels", "coex_restrictions", "wifi6_models"
         ]
         self.unpack_userparams(
             req_param_names=req_params, opt_param_names=opt_param)
@@ -106,9 +106,8 @@
         self.turn_location_off_and_scan_toggle_off()
         if self.dut.droid.wifiIsApEnabled():
             wutils.stop_wifi_tethering(self.dut)
-        wutils.reset_wifi(self.dut)
-        if self.dut_client:
-            wutils.reset_wifi(self.dut_client)
+        for ad in self.android_devices:
+            wutils.reset_wifi(ad)
 
     def teardown_class(self):
         if "AccessPoint" in self.user_params:
@@ -242,6 +241,8 @@
                        (network_ssid, connect_ssid))
         if connect_ssid != network_ssid:
             return False
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
         return True
 
     def run_iperf_client(self, params):
@@ -385,6 +386,8 @@
         if not reconnect:
             raise signals.TestFailure("Device did not connect to the correct"
                                       " network after toggling WiFi.")
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def helper_reconnect_toggle_airplane(self):
         """Connect to multiple networks, turn on/off Airplane moce, then
@@ -408,6 +411,8 @@
         if not reconnect:
             raise signals.TestFailure("Device did not connect to the correct"
                                       " network after toggling Airplane mode.")
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def helper_reboot_configstore_reconnect(self, lock_screen=False):
         """Connect to multiple networks, reboot then reconnect to previously
@@ -439,6 +444,8 @@
             raise signals.TestFailure(
                 "Device failed to reconnect to the correct"
                 " network after reboot.")
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def helper_toggle_wifi_reboot_configstore_reconnect(self):
         """Connect to multiple networks, disable WiFi, reboot, then
@@ -473,6 +480,8 @@
             msg = ("Device failed to reconnect to the correct network after"
                    " toggling WiFi and rebooting.")
             raise signals.TestFailure(msg)
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def helper_toggle_airplane_reboot_configstore_reconnect(self):
         """Connect to multiple networks, enable Airplane mode, reboot, then
@@ -511,6 +520,8 @@
             msg = ("Device failed to reconnect to the correct network after"
                    " toggling Airplane mode and rebooting.")
             raise signals.TestFailure(msg)
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def verify_traffic_between_devices(self,dest_device,src_device,num_of_tries=2):
         """Test the clients and DUT can ping each other.
@@ -862,16 +873,22 @@
         self.dut.droid.wakeLockRelease()
         self.dut.droid.goToSleepNow()
         wutils.connect_to_wifi_network(self.dut, self.wpapsk_5g)
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="81eb7527-4c92-4422-897a-6b5f6445e84a")
     def test_config_store_with_wpapsk_2g(self):
         self.connect_to_wifi_network_toggle_wifi_and_run_iperf(
             (self.wpapsk_2g, self.dut))
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="8457903d-cb7e-4c89-bcea-7f59585ea6e0")
     def test_config_store_with_wpapsk_5g(self):
         self.connect_to_wifi_network_toggle_wifi_and_run_iperf(
             (self.wpapsk_5g, self.dut))
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="b9fbc13a-47b4-4f64-bd2c-e5a3cb24ab2f")
     def test_tdls_supported(self):
@@ -899,6 +916,8 @@
         """
         wutils.wifi_connect(self.dut, self.open_network_2g)
         self.get_energy_info()
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="2622c253-defc-4a35-93a6-ca9d29a8238c")
     def test_connect_to_wep_2g(self):
@@ -929,6 +948,8 @@
         2. Connect to the network and validate internet connection.
         """
         wutils.connect_to_wifi_network(self.dut, self.wpa_networks[0]["2g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="612c3c31-a4c5-4014-9a2d-3f4bcc20c0d7")
     def test_connect_to_wpa_5g(self):
@@ -939,6 +960,8 @@
         2. Connect to the network and validate internet connection.
         """
         wutils.connect_to_wifi_network(self.dut, self.wpa_networks[0]["5g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     @test_tracker_info(uuid="2a617fb4-1d8e-44e9-a500-a5456e1df83f")
     def test_connect_to_2g_can_be_pinged(self):
@@ -950,7 +973,11 @@
         3. Check DUT can be pinged by another device
         """
         wutils.connect_to_wifi_network(self.dut, self.wpa_networks[0]["2g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
         wutils.connect_to_wifi_network(self.dut_client, self.wpa_networks[0]["2g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut_client, self.wifi6_models, "wifi6_ap" in self.user_params)
         self.verify_traffic_between_devices(self.dut,self.dut_client)
 
     @test_tracker_info(uuid="94bdd657-649b-4a2c-89c3-3ec6ba18e08e")
@@ -963,7 +990,11 @@
         3. Check DUT can be pinged by another device
         """
         wutils.connect_to_wifi_network(self.dut, self.wpa_networks[0]["5g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
         wutils.connect_to_wifi_network(self.dut_client, self.wpa_networks[0]["5g"])
+        wutils.verify_11ax_wifi_connection(
+            self.dut_client, self.wifi6_models, "wifi6_ap" in self.user_params)
         self.verify_traffic_between_devices(self.dut,self.dut_client)
 
     @test_tracker_info(uuid="d87359aa-c4da-4554-b5de-8e3fa852a6b0")
@@ -1084,6 +1115,8 @@
         """
         network = self.open_network_5g
         wutils.connect_to_wifi_network(self.dut, network)
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
         info = self.dut.droid.wifiGetConnectionInfo()
         network_id = info[WifiEnums.NETID_KEY]
         self.dut.log.info("Disable auto join on network")
@@ -1095,6 +1128,8 @@
                                     assert_on_fail=False), "Device should not connect.")
         self.dut.droid.wifiEnableAutojoin(network_id, True)
         wutils.wait_for_connect(self.dut, network[WifiEnums.SSID_KEY], assert_on_fail=False)
+        wutils.verify_11ax_wifi_connection(
+            self.dut, self.wifi6_models, "wifi6_ap" in self.user_params)
 
     def coex_unsafe_channel_key(self, unsafe_channel):
         if COEX_POWER_CAP_DBM in unsafe_channel:
diff --git a/acts_tests/tests/google/wifi/wifi6/WifiManager11axTest.py b/acts_tests/tests/google/wifi/wifi6/WifiManager11axTest.py
new file mode 100644
index 0000000..d2fb981
--- /dev/null
+++ b/acts_tests/tests/google/wifi/wifi6/WifiManager11axTest.py
@@ -0,0 +1,84 @@
+#
+#   Copyright 2021 - The Android Open Source Project
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from acts_contrib.test_utils.wifi.WifiBaseTest import WifiBaseTest
+from WifiManagerTest import WifiManagerTest
+
+
+class WifiManager11axTest(WifiManagerTest):
+  """Tests for WifiManager 11ax.
+
+  Test Bed Requirement:
+    One Android device and 2 Asus AXE11000 Access Point.
+  """
+
+  def __init__(self, configs):
+    super().__init__(configs)
+    self.tests = (
+        "test_toggle_wifi_state_and_get_startupTime",
+        "test_toggle_with_screen",
+        "test_scan",
+        "test_scan_with_wifi_off_and_location_scan_on",
+        "test_scan_after_reboot_with_wifi_off_and_location_scan_on",
+        "test_scan_with_wifi_off_and_location_scan_off",
+        "test_add_network",
+        "test_forget_network",
+        "test_reconnect_to_previously_connected",
+        "test_reconnect_toggle_wifi",
+        "test_reconnect_toggle_wifi_with_location_scan_on",
+        "test_reconnect_toggle_airplane",
+        "test_reconnect_toggle_airplane_with_location_scan_on",
+        "test_reboot_configstore_reconnect",
+        "test_reboot_configstore_reconnect_with_location_scan_on",
+        "test_toggle_wifi_reboot_configstore_reconnect",
+        "test_toggle_wifi_reboot_configstore_reconnect_with_location_scan_on",
+        "test_toggle_airplane_reboot_configstore_reconnect",
+        "test_toggle_airplane_reboot_configstore_reconnect_with_location_scan_on",
+        "test_reboot_configstore_reconnect_with_screen_lock",
+        "test_connect_to_5g_after_reboot_without_unlock",
+        "test_config_store_with_wpapsk_2g",
+        "test_config_store_with_wpapsk_5g",
+        "test_tdls_supported",
+        "test_energy_info",
+        "test_energy_info_connected",
+        "test_connect_to_wpa_2g",
+        "test_connect_to_wpa_5g",
+        "test_connect_to_2g_can_be_pinged",
+        "test_connect_to_5g_can_be_pinged",
+        "test_wifi_saved_network_reset",
+        "test_reboot_wifi_and_bluetooth_on",
+        "test_scan_result_api",
+        "test_enable_disable_auto_join_saved_network",
+        "test_set_get_coex_unsafe_channels"
+    )
+
+  def setup_class(self):
+    WifiBaseTest.setup_class(self)
+
+    self.dut = self.android_devices[0]
+    self.dut_client = self.android_devices[1]
+    req_params = [
+        "reference_networks", "wpa_networks", "iperf_server_address",
+        "iperf_server_port", "coex_unsafe_channels", "coex_restrictions",
+        "wifi6_models"
+    ]
+    self.unpack_userparams(req_param_names=req_params,)
+    self.ap = self.access_points[1]
+    self.ap.configure_ap({"2g": {"security": "open"},
+                          "5g": {"security": "open"}})
+    self.wpapsk_2g = self.reference_networks[0]["2g"]
+    self.wpapsk_5g = self.reference_networks[0]["5g"]
+    self.open_network_2g = self.ap.get_wifi_network("2g")
+    self.open_network_5g = self.ap.get_wifi_network("5g")