[ATFT] Bug fix: mapped device not shown correctly.android-q-preview-6android-q-preview-5android-q-preview-4android-q-preview-2.5android-q-preview-1

Fix a bug that would cause mapped device to be hidden when the
mapped device location is not in ascending order. For example, slot
1 is mapped to device location 2, and slot 2 is mapped to device
location 1.

Tests: Added unit test testPrintTargetDevicesMultipleDeviceModeReverseOrder
to cover this case. Failed before the change and succeed afterwards.
Bug: b/123679718

2 files changed, 60 insertions, 14 deletions

diff --git a/at-factory-tool/atft.py b/at-factory-tool/atft.py
index e0fc098..27dfdd3 100644
--- a/at-factory-tool/atft.py
+++ b/at-factory-tool/atft.py
@@ -3746,7 +3746,6 @@ class Atft(wx.Frame):
           self.unmapped_target_dev_component, serial_number, serial_text,
           status, state)
     else:
-      target_dev_index = 0
       for i in range(TARGET_DEV_SIZE):
         serial_text = ''
         status = None
@@ -3755,15 +3754,13 @@ class Atft(wx.Frame):
         target_dev_component = self.target_devs_components[i]
         if self.device_usb_locations[i]:
           target_dev_component.active = True
-          if (target_dev_index < len(target_devs) and
-              (target_devs[target_dev_index].location ==
-               self.device_usb_locations[i])):
-            serial_number = target_devs[target_dev_index].serial_number
-            serial_text = '{}: {}'.format(
-                self.atft_string.FIELD_SERIAL_NUMBER, str(serial_number))
-            status = target_devs[target_dev_index].provision_status
-            state = target_devs[target_dev_index].provision_state
-            target_dev_index += 1
+          for target_dev in target_devs:
+            if (target_dev.location == self.device_usb_locations[i]):
+              serial_number = target_dev.serial_number
+              serial_text = '{}: {}'.format(
+                  self.atft_string.FIELD_SERIAL_NUMBER, str(serial_number))
+              status = target_dev.provision_status
+              state = target_dev.provision_state
         else:
           target_dev_component.active = False
         self._ShowTargetDevice(
diff --git a/at-factory-tool/atft_unittest.py b/at-factory-tool/atft_unittest.py
index 40f160a..576749f 100644
--- a/at-factory-tool/atft_unittest.py
+++ b/at-factory-tool/atft_unittest.py
@@ -66,6 +66,7 @@ class MockAtft(atft.Atft):
     self.target_devs_components = MagicMock()
     atft.Atft.__init__(self)
     self.provision_steps = self.DEFAULT_PROVISION_STEPS_PRODUCT
+    self.skip_reboot = False
 
   def _MockParseConfig(self):
     self.atft_version = 'vTest'
@@ -225,14 +226,14 @@ class AtftTest(TestCase):
     mock_atft._CheckMappingMode.assert_not_called()
 
   def testPrintTargetDevicesMultipleDeviceMode(self):
-      # Test _PrintTargetDevices in multiple device mode.
+    # Test _PrintTargetDevices in multiple device mode.
     mock_atft = MockAtft()
     mock_serial_string = MagicMock()
     mock_atft.atft_string.FIELD_SERIAL_NUMBER = mock_serial_string
     mock_atft.atft_manager = MagicMock()
     mock_dev_components = []
     for i in range(0, 6):
-        mock_dev_components.append(MagicMock())
+      mock_dev_components.append(MagicMock())
     mock_atft.target_devs_components = mock_dev_components;
     dev1 = self.test_dev1
     dev2 = self.test_dev2
@@ -270,6 +271,54 @@ class AtftTest(TestCase):
             ProvisionStatus.REBOOT_IN_PROGRESS, dev2_state)
         ])
 
+  def testPrintTargetDevicesMultipleDeviceModeReverseOrder(self):
+    # Test _PrintTargetDevices in multiple device mode with the mapped device
+    # location in descending order. This test is to verify a bug that was caused
+    # by the assumption that the mapped device location is in ascending order.
+    mock_atft = MockAtft()
+    mock_serial_string = MagicMock()
+    mock_atft.atft_string.FIELD_SERIAL_NUMBER = mock_serial_string
+    mock_atft.atft_manager = MagicMock()
+    mock_dev_components = []
+    for i in range(0, 6):
+      mock_dev_components.append(MagicMock())
+    mock_atft.target_devs_components = mock_dev_components;
+    dev1 = self.test_dev1
+    dev2 = self.test_dev2
+    dev1.provision_status = ProvisionStatus.IDLE
+    dev2.provision_status = ProvisionStatus.REBOOT_IN_PROGRESS
+    dev1_state = ProvisionState()
+    dev1.provision_state = dev1_state
+    dev2_state = ProvisionState()
+    dev2_state.bootloader_locked = True
+    dev2_state.avb_perm_attr_set = True
+    dev2_state.avb_locked = True
+    dev2.provision_state = dev2_state
+    mock_atft.atft_manager.target_devs = [dev1, dev2]
+    mock_atft.device_usb_locations = []
+    for i in range(mock_atft.TARGET_DEV_SIZE):
+      mock_atft.device_usb_locations.append(None)
+    # Two target devices at location 0 and location 5.
+    mock_atft.device_usb_locations[0] = self.TEST_LOCATION2
+    mock_atft.device_usb_locations[5] = self.TEST_LOCATION1
+    mock_atft._ShowTargetDevice = MagicMock()
+    mock_atft._PrintTargetDevices()
+    mock_atft._ShowTargetDevice.assert_has_calls([
+        call(
+            mock_dev_components[0],
+            self.TEST_SERIAL2,
+            '{}: {}'.format(mock_serial_string, self.TEST_SERIAL2),
+            ProvisionStatus.REBOOT_IN_PROGRESS, dev2_state),
+        call(mock_dev_components[1], None, '', None, None),
+        call(mock_dev_components[2], None, '', None, None),
+        call(mock_dev_components[3], None, '', None, None),
+        call(mock_dev_components[4], None, '', None, None),
+        call(
+            mock_dev_components[5], self.TEST_SERIAL1,
+            '{}: {}'.format(mock_serial_string, self.TEST_SERIAL1),
+            ProvisionStatus.IDLE, dev1_state)
+        ])
+
   def testPrintTargetDevicesSingleDeviceMode(self):
     # Test _PrintTargetDevices in single device mode.
     mock_atft = MockAtft()
@@ -1256,11 +1305,11 @@ class AtftTest(TestCase):
   def MockGetTargetDevice(self, serial):
     return self.device_map.get(serial)
 
-  def MockReboot(self, target, timeout, success, fail):
+  def MockReboot(self, target, timeout, success, fail, skip_reboot):
     success()
     target.provision_state.bootloader_locked = True
 
-  def MockRebootStateNoChange(self, target, timeout, success, fail):
+  def MockRebootStateNoChange(self, target, timeout, success, fail, skip_reboot):
     success()
     target.provision_state.bootloader_locked = False