Fix bug in wifi sensitivity tests

This CL fixes a bug in wifi sensitivity tests where attenuator offsets
were not set properly. Further the CL improves test timing by not
attempting to detect the attenuator to DUT chain mapping if it is
already known from a previous test on the same channel.

Test: Done
Bug: None

Change-Id: I3b5d1dc3be2326f9134259fb391e5bd37e11db94
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 a982cc6..5b5c298 100644
--- a/acts/tests/google/wifi/WifiSensitivityTest.py
+++ b/acts/tests/google/wifi/WifiSensitivityTest.py
@@ -173,14 +173,14 @@
             data_msc = [file for file in self.firmware
                         if "Data.msc" in file][0]
             wputils.push_firmware(self.dut, wlanmdsp, data_msc)
+        self.atten_dut_chain_map = {}
         self.testclass_results = []
 
         # Turn WiFi ON
         if self.testclass_params.get('airplane_mode', 1):
             self.log.info('Turning on airplane mode.')
-            asserts.assert_true(
-                utils.force_airplane_mode(self.dut, True),
-                "Can not turn on airplane mode.")
+            asserts.assert_true(utils.force_airplane_mode(self.dut, True),
+                                "Can not turn on airplane mode.")
         wutils.wifi_toggle_state(self.dut, True)
 
     def teardown_class(self):
@@ -356,9 +356,10 @@
         self.access_point.set_channel(band, testcase_params['channel'])
         self.access_point.set_bandwidth(band, testcase_params['mode'])
         self.access_point.set_power(band, testcase_params['ap_tx_power'])
-        self.access_point.set_rate(
-            band, testcase_params['mode'], testcase_params['num_streams'],
-            testcase_params['rate'], testcase_params['short_gi'])
+        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'])]
@@ -367,6 +368,8 @@
                 atten.offset = atten_offsets[0]
             elif 'AP-Chain-1' in atten.path:
                 atten.offset = atten_offsets[1]
+            else:
+                atten.offset = 0
         self.log.info('Access Point Configuration: {}'.format(
             self.access_point.ap_settings))
 
@@ -396,18 +399,20 @@
             self.dut.droid.wifiSetCountryCode(
                 self.testclass_params['country_code'])
             self.main_network[band]['channel'] = testcase_params['channel']
-            wutils.wifi_connect(
-                self.dut,
-                self.main_network[band],
-                num_of_tries=5,
-                check_connectivity=False)
+            wutils.wifi_connect(self.dut,
+                                self.main_network[band],
+                                num_of_tries=5,
+                                check_connectivity=False)
         self.dut_ip = self.dut.droid.connectivityGetIPv4Addresses('wlan0')[0]
-        atten_dut_chain_map = wputils.get_current_atten_dut_chain_map(
-            self.attenuators, self.dut, self.ping_server)
-        self.log.info(
-            "Current Attenuator-DUT Chain Map: {}".format(atten_dut_chain_map))
+        self.atten_dut_chain_map.setdefault(
+            testcase_params['channel'],
+            wputils.get_current_atten_dut_chain_map(self.attenuators, self.dut,
+                                                    self.ping_server))
+        self.log.info("Current Attenuator-DUT Chain Map: {}".format(
+            self.atten_dut_chain_map[testcase_params['channel']]))
         for idx, atten in enumerate(self.attenuators):
-            if atten_dut_chain_map[idx] == testcase_params['attenuated_chain']:
+            if self.atten_dut_chain_map[testcase_params['channel']][
+                    idx] == testcase_params['attenuated_chain']:
                 atten.offset = atten.instrument.max_atten
 
     def extract_test_id(self, testcase_params, id_fields):
@@ -578,8 +583,9 @@
 class WifiSensitivity_SampleChannels_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         super().__init__(controllers)
-        self.tests = self.generate_test_cases(
-            [6, 36, 149], ['VHT20', 'VHT40', 'VHT80'], ['0', '1', '2x2'])
+        self.tests = self.generate_test_cases([6, 36, 149],
+                                              ['VHT20', 'VHT40', 'VHT80'],
+                                              ['0', '1', '2x2'])
 
 
 class WifiSensitivity_2GHz_Test(WifiSensitivityTest):
@@ -600,8 +606,9 @@
 class WifiSensitivity_UNII1_Test(WifiSensitivityTest):
     def __init__(self, controllers):
         super().__init__(controllers)
-        self.tests = self.generate_test_cases(
-            [36, 40, 44, 48], ['VHT20', 'VHT40', 'VHT80'], ['0', '1', '2x2'])
+        self.tests = self.generate_test_cases([36, 40, 44, 48],
+                                              ['VHT20', 'VHT40', 'VHT80'],
+                                              ['0', '1', '2x2'])
 
 
 class WifiSensitivity_UNII3_Test(WifiSensitivityTest):
@@ -620,7 +627,6 @@
     It allows setting orientation and other chamber parameters to study
     performance in varying channel conditions
     """
-
     def __init__(self, controllers):
         base_test.BaseTestClass.__init__(self, controllers)
         self.testcase_metric_logger = (
@@ -691,11 +697,10 @@
                 x_label='Orientation (deg)',
                 primary_y_label='Sensitivity (dBm)')
             for channel, channel_results in test_data.items():
-                curr_plot.add_line(
-                    channel_results['orientation'],
-                    channel_results['sensitivity'],
-                    legend='Channel {}'.format(channel),
-                    marker='circle')
+                curr_plot.add_line(channel_results['orientation'],
+                                   channel_results['sensitivity'],
+                                   legend='Channel {}'.format(channel),
+                                   marker='circle')
                 metric_tag = 'ota_summary_ch{}_{}'.format(
                     channel, metric_test_config)
                 metric_name = metric_tag + '.avg_sensitivity'
@@ -808,9 +813,10 @@
             self.RateTuple(8, 2, 173.3),
             self.RateTuple(2, 2, 43.3)
         ]
-        self.tests = self.generate_test_cases(
-            requested_channels, ['VHT20', 'VHT80'], requested_rates, ['2x2'],
-            list(range(0, 360, 10)))
+        self.tests = self.generate_test_cases(requested_channels,
+                                              ['VHT20', 'VHT80'],
+                                              requested_rates, ['2x2'],
+                                              list(range(0, 360, 10)))
 
 
 class WifiOtaSensitivity_SingleChain_TenDegree_Test(WifiOtaSensitivityTest):
@@ -821,9 +827,10 @@
             self.RateTuple(8, 1, 86.7),
             self.RateTuple(2, 1, 21.7)
         ]
-        self.tests = self.generate_test_cases(
-            requested_channels, ['VHT20', 'VHT80'], requested_rates, ['2x2'],
-            list(range(0, 360, 10)))
+        self.tests = self.generate_test_cases(requested_channels,
+                                              ['VHT20', 'VHT80'],
+                                              requested_rates, ['2x2'],
+                                              list(range(0, 360, 10)))
 
 
 class WifiOtaSensitivity_ThirtyDegree_Test(WifiOtaSensitivityTest):
@@ -844,9 +851,10 @@
             self.RateTuple(2, 2, 43.3),
             self.RateTuple(0, 2, 14.4)
         ]
-        self.tests = self.generate_test_cases(
-            requested_channels, ['VHT20', 'VHT80'], requested_rates, ['2x2'],
-            list(range(0, 360, 30)))
+        self.tests = self.generate_test_cases(requested_channels,
+                                              ['VHT20', 'VHT80'],
+                                              requested_rates, ['2x2'],
+                                              list(range(0, 360, 30)))
 
 
 class WifiOtaSensitivity_45Degree_Test(WifiOtaSensitivityTest):