Antenna verification test for router

Added tests to verify each antenna in the router is working.

BUG=chromium:343184
TEST=1. Verify no connection failure when running this test.
     2. To induce a "failure" with this test, manually put a terminator on one antenna in the router.
	Run this test.
	Verify connection failure or connected with weak signal (signal leak from terminator) when the terminated antenna is the only one active

Change-Id: I47e9e01ee6489b013b29d8f36238da9d1e425268
Reviewed-on: https://chromium-review.googlesource.com/187337
Reviewed-by: Christopher Wiley <wiley@chromium.org>
Commit-Queue: Peter Qiu <zqiu@chromium.org>
Tested-by: Peter Qiu <zqiu@chromium.org>
diff --git a/client/common_lib/cros/network/iw_runner.py b/client/common_lib/cros/network/iw_runner.py
index 024599f..2def72f 100644
--- a/client/common_lib/cros/network/iw_runner.py
+++ b/client/common_lib/cros/network/iw_runner.py
@@ -42,7 +42,8 @@
 #          "authenticate".
 #   max_scan_ssids: Maximum number of SSIDs which can be scanned at once.
 IwPhy = collections.namedtuple(
-    'Phy', ['name', 'bands', 'modes', 'commands', 'max_scan_ssids'])
+    'Phy', ['name', 'bands', 'modes', 'commands', 'max_scan_ssids',
+            'avail_tx_antennas', 'avail_rx_antennas'])
 
 DEFAULT_COMMAND_IW = 'iw'
 
@@ -278,7 +279,9 @@
                             bands,
                             tuple(pending_phy_modes),
                             tuple(pending_phy_commands),
-                            pending_phy_max_scan_ssids)
+                            pending_phy_max_scan_ssids,
+                            pending_phy_tx_antennas,
+                            pending_phy_rx_antennas)
             all_phys.append(new_phy)
 
         for line in output.splitlines():
@@ -291,6 +294,8 @@
                 pending_phy_modes = []
                 pending_phy_commands = []
                 pending_phy_max_scan_ssids = None
+                pending_phy_tx_antennas = 0
+                pending_phy_rx_antennas = 0
                 continue
 
             match_section = re.match('\s*(\w.*):\s*$', line)
@@ -326,6 +331,15 @@
                     pending_phy_commands.append(command_match.group(1))
                     continue
 
+            match_avail_antennas = re.match('\s*Available Antennas: TX (\S+)'
+                                            ' RX (\S+)', line)
+            if match_avail_antennas and pending_phy_name:
+                pending_phy_tx_antennas = int(
+                        match_avail_antennas.group(1), 16)
+                pending_phy_rx_antennas = int(
+                        match_avail_antennas.group(2), 16)
+                continue
+
             if not all([current_band, pending_phy_name,
                         line.startswith('\t')]):
                 continue
@@ -523,3 +537,20 @@
                 return True
             time.sleep(1)
         return False
+
+
+    def set_antenna_bitmap(self, phy, tx_bitmap, rx_bitmap):
+        """Set antenna chain mask on given phy (radio).
+
+        This function will set the antennas allowed to use for TX and
+        RX on the |phy| based on the |tx_bitmap| and |rx_bitmap|.
+        This command is only allowed when the interfaces on the phy are down.
+
+        @param phy: phy name
+        @param tx_bitmap: bitmap of allowed antennas to use for TX
+        @param rx_bitmap: bitmap of allowed antennas to use for RX
+
+        """
+        command = '%s phy %s set antenna %d %d' % (self._command_iw, phy,
+                                                   tx_bitmap, rx_bitmap)
+        self._run(command)
diff --git a/server/site_linux_router.py b/server/site_linux_router.py
index d97e033..4e64a8c 100644
--- a/server/site_linux_router.py
+++ b/server/site_linux_router.py
@@ -111,6 +111,9 @@
         # Place us in the US by default
         self.iw_runner.set_regulatory_domain('US')
 
+        # Reset all antennas to be active
+        self.set_default_antenna_bitmap()
+
 
     def close(self):
         """Close global resources held by this system."""
diff --git a/server/site_linux_system.py b/server/site_linux_system.py
index eb07e74..0c3b156 100644
--- a/server/site_linux_system.py
+++ b/server/site_linux_system.py
@@ -337,3 +337,21 @@
         if missing:
             raise to_be_raised('AP on %s is missing required capabilites: %r' %
                                (self.role, missing))
+
+
+    def set_antenna_bitmap(self, tx_bitmap, rx_bitmap):
+        """Setup antenna bitmaps for all the phys.
+
+        @param tx_bitmap int bitmap of allowed antennas to use for TX
+        @param rx_bitmap int bitmap of allowed antennas to use for RX
+
+        """
+        for phy in self.phy_list:
+            self.iw_runner.set_antenna_bitmap(phy.name, tx_bitmap, rx_bitmap)
+
+
+    def set_default_antenna_bitmap(self):
+        """Setup default antenna bitmaps for all the phys."""
+        for phy in self.phy_list:
+            self.iw_runner.set_antenna_bitmap(phy.name, phy.avail_tx_antennas,
+                                              phy.avail_rx_antennas)
diff --git a/server/site_tests/network_WiFi_VerifyRouter/network_WiFi_VerifyRouter.py b/server/site_tests/network_WiFi_VerifyRouter/network_WiFi_VerifyRouter.py
index 1782149..c237ee3 100644
--- a/server/site_tests/network_WiFi_VerifyRouter/network_WiFi_VerifyRouter.py
+++ b/server/site_tests/network_WiFi_VerifyRouter/network_WiFi_VerifyRouter.py
@@ -14,6 +14,46 @@
     """Test that a dual radio router can use both radios."""
     version = 1
 
+    def _antenna_test(self, channel):
+        """Test to verify each antenna is working on given band.
+
+        Setup AP in each radio in given band, and run connection test with one
+        antenna active at a time, to verify each antenna in each radio is
+        working correctly for given band. Antenna can only be configured when
+        the wireless interface is down.
+
+        @param channel: int Wifi channel to conduct test on
+
+        """
+        # Connect to AP with only one antenna active at a time
+        for bitmap in (1,2):
+            self.context.router.deconfig()
+            self.context.router.set_antenna_bitmap(bitmap, bitmap)
+            # Setup two APs in the same band
+            n_mode = hostap_config.HostapConfig.MODE_11N_MIXED
+            ap_config = hostap_config.HostapConfig(channel=channel, mode=n_mode)
+            self.context.configure(ap_config)
+            self.context.configure(ap_config, multi_interface=True)
+            # Verify connectivity to both APs
+            for instance in range(2):
+                client_conf = xmlrpc_datatypes.AssociationParameters(
+                        ssid=self.context.router.get_ssid(instance=instance))
+                self.context.assert_connect_wifi(client_conf)
+                logging.info('Signal level for AP %d with bitmap %d is %d',
+                             instance, bitmap,
+                             self.context.client.wifi_signal_level)
+
+
+    def cleanup(self):
+        """Clean up after the test is completed
+
+        Perform additional cleanups after the test, the important thing is
+        to restore default antennas bitmap.
+        """
+        self.context.router.deconfig()
+        self.context.router.set_default_antenna_bitmap()
+        super(network_WiFi_VerifyRouter, self).cleanup()
+
 
     def run_once(self):
         """Set up two APs connect to both and then exit."""
@@ -31,3 +71,7 @@
                          self.context.client.wifi_signal_level)
             self.write_perf_keyval({'signal_for_ap_%d' % instance:
                                     self.context.client.wifi_signal_level})
+
+        # Run antenna test for 2GHz band and 5GHz band
+        self._antenna_test(6)
+        self._antenna_test(136)