Add single chain senstiivity tests

This CL adds the capability to do single chaing sensitivity tests by
attenuating the chain-not-in-test up to the maximum attenuation
supported by the attenuator. The CL also includes additional results
generation in a CSV file for better readability. To achieve the single
chain tests, the CL includes changes to the attenuator class to be more
permissive when setting attenuation levels above the maximum supported;
this is allowed via an optional flag passed to set_atten which defaults
to strict attenuation checking, i.e., current behavior.

Test: Done
Bug: None

Change-Id: I750a8abbb5867e22acfe215f82ba42b9bcd76207
Signed-off-by: Omar El Ayach <oelayach@google.com>
diff --git a/acts/tests/google/wifi/WifiSensitivityTest.py b/acts/tests/google/wifi/WifiSensitivityTest.py
index 66eb0b5..9c554e4 100644
--- a/acts/tests/google/wifi/WifiSensitivityTest.py
+++ b/acts/tests/google/wifi/WifiSensitivityTest.py
@@ -15,6 +15,8 @@
 #   limitations under the License.
 
 import collections
+import csv
+import itertools
 import json
 import logging
 import os
@@ -104,6 +106,12 @@
         for dev in self.android_devices:
             wutils.wifi_toggle_state(dev, True)
 
+    def teardown_class(self):
+        # Turn WiFi OFF
+        for dev in self.android_devices:
+            wutils.wifi_toggle_state(dev, False)
+        self.process_testclass_results()
+
     def pass_fail_check(self, result):
         """Checks sensitivity against golden results and decides on pass/fail.
 
@@ -132,6 +140,7 @@
 
     def process_testclass_results(self):
         """Saves and plots test results from all executed test cases."""
+        # write json output
         testclass_results_dict = collections.OrderedDict()
         for result in self.testclass_results:
             testclass_results_dict[result["test_name"]] = {
@@ -142,6 +151,28 @@
         results_file_path = os.path.join(self.log_path, 'results.json')
         with open(results_file_path, 'w') as results_file:
             json.dump(testclass_results_dict, results_file, indent=4)
+        # write csv
+        results_file_path = os.path.join(self.log_path, 'results.csv')
+        with open(results_file_path, mode='w') as csv_file:
+            csv_header = [
+                "Channel", "Mode", "MCS", "Streams", "Chain", "Sensitivity",
+                "Range", "Peak Throughput"
+            ]
+            writer = csv.DictWriter(csv_file, fieldnames=csv_header)
+            writer.writeheader()
+            for result in self.testclass_results:
+                testcase_params = self.parse_test_params(result["test_name"])
+                writer.writerow({
+                    "Channel": testcase_params["channel"],
+                    "Mode": testcase_params["mode"],
+                    "MCS": testcase_params["rate"],
+                    "Streams": testcase_params["num_streams"],
+                    "Chain": testcase_params["chain_mask"],
+                    "Sensitivity": result["sensitivity"],
+                    "Range": result["range"],
+                    "Peak Throughput": result["peak_throughput"]
+                })
+
         if not self.testclass_params["traffic_type"].lower() == "ping":
             WifiRvrTest.process_testclass_results(self)
 
@@ -219,13 +250,16 @@
         self.access_point.set_rate(
             band, testcase_params["mode"], testcase_params["num_streams"],
             testcase_params["rate"], testcase_params["short_gi"])
+        # Set attenuator offsets and set attenuators to initial condition
+        atten_offsets = self.testbed_params['chain_offset'][str(
+            testcase_params['channel'])]
         for atten in self.attenuators:
-            if atten.path == "Chain-0":
-                atten.offset = self.testbed_params["chain_offset"][str(
-                    testcase_params["channel"])][0]
-            elif atten.path == "Chain-1":
-                atten.offset = self.testbed_params["chain_offset"][str(
-                    testcase_params["channel"])][1]
+            if 'AP-Chain-0' in atten.path:
+                atten.offset = atten_offsets[0]
+            elif 'AP-Chain-1' in atten.path:
+                atten.offset = atten_offsets[1]
+            if testcase_params["attenuated_chain"] in atten.path:
+                atten.offset = atten.instrument.max_atten
         self.log.info("Access Point Configuration: {}".format(
             self.access_point.ap_settings))
 
@@ -286,6 +320,12 @@
             testcase_params["rate"] = int(test_name_params[4][3:])
         testcase_params["num_streams"] = int(test_name_params[5][3:])
         testcase_params["short_gi"] = 0
+        testcase_params["chain_mask"] = test_name_params[6][2:]
+        if testcase_params["chain_mask"] in ["0", "1"]:
+            testcase_params["attenuated_chain"] = "DUT-Chain-{}".format(
+                1 if testcase_params['chain_mask'] == "0" else 0)
+        else:
+            testcase_params["attenuated_chain"] = None
 
         if self.testclass_params["traffic_type"] == "UDP":
             testcase_params["iperf_args"] = '-i 1 -t {} -J -u -b {} -R'.format(
@@ -330,7 +370,7 @@
         self.testclass_results.append(result)
         self.pass_fail_check(result)
 
-    def generate_test_cases(self, channels):
+    def generate_test_cases(self, channels, chain_mask):
         """Function that auto-generates test cases for a test class."""
         testcase_wrapper = self._test_sensitivity
         for channel in channels:
@@ -345,14 +385,17 @@
                     rates = self.VALID_RATES["legacy_5GHz"]
                 else:
                     raise ValueError("Invalid test mode.")
-                for rate in rates:
+                for chain, rate in itertools.product(chain_mask, rates):
+                    if str(chain) in ["0", "1"] and rate[1] == 2:
+                        # Do not test 2-stream rates in single chain mode
+                        continue
                     if "legacy" in mode:
-                        testcase_name = "test_sensitivity_ch{}_{}_{}_nss{}".format(
+                        testcase_name = "test_sensitivity_ch{}_{}_{}_nss{}_ch{}".format(
                             channel, mode,
-                            str(rate[0]).replace(".", "p"), rate[1])
+                            str(rate[0]).replace(".", "p"), rate[1], chain)
                     else:
-                        testcase_name = "test_sensitivity_ch{}_{}_mcs{}_nss{}".format(
-                            channel, mode, rate[0], rate[1])
+                        testcase_name = "test_sensitivity_ch{}_{}_mcs{}_nss{}_ch{}".format(
+                            channel, mode, rate[0], rate[1], chain)
                     setattr(self, testcase_name, testcase_wrapper)
                     self.tests.append(testcase_name)
 
@@ -361,106 +404,108 @@
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
         self.generate_test_cases(
-            [1, 2, 6, 10, 11, 36, 40, 44, 48, 149, 153, 157, 161])
+            [1, 2, 6, 10, 11, 36, 40, 44, 48, 149, 153, 157, 161],
+            ["0", "1", "2x2"])
 
 
 class WifiSensitivity_2GHz_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([1, 2, 6, 10, 11])
+        self.generate_test_cases([1, 2, 6, 10, 11], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_5GHz_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([36, 40, 44, 48, 149, 153, 157, 161])
+        self.generate_test_cases([36, 40, 44, 48, 149, 153, 157, 161],
+                                 ["0", "1", "2x2"])
 
 
 class WifiSensitivity_UNII1_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([36, 40, 44, 48])
+        self.generate_test_cases([36, 40, 44, 48], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_UNII3_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([149, 153, 157, 161])
+        self.generate_test_cases([149, 153, 157, 161], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch1_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([1])
+        self.generate_test_cases([1], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch2_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([2])
+        self.generate_test_cases([2], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch6_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([6])
+        self.generate_test_cases([6], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch10_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([10])
+        self.generate_test_cases([10], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch11_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([11])
+        self.generate_test_cases([11], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch36_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([36])
+        self.generate_test_cases([36], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch40_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([40])
+        self.generate_test_cases([40], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch44_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([44])
+        self.generate_test_cases([44], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch48_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([48])
+        self.generate_test_cases([48], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch149_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([149])
+        self.generate_test_cases([149], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch153_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([153])
+        self.generate_test_cases([153], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch157_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([157])
+        self.generate_test_cases([157], ["0", "1", "2x2"])
 
 
 class WifiSensitivity_ch161_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
-        self.generate_test_cases([161])
+        self.generate_test_cases([161], ["0", "1", "2x2"])