Add more test cases for get_component_devs

repos/system_upgrade/common/libraries/tests/test_grub.py

@@ -12,7 +12,7 @@
 BOOT_DEVICE = '/dev/vda'
 
 MD_BOOT_DEVICE = '/dev/md0'
-MD_BOOT_REAL_DEVICES = ['/dev/sda', '/dev/sdb']
+MD_BOOT_DEVICES_WITH_GRUB = ['/dev/sda', '/dev/sdb']
 
 VALID_DD = b'GRUB GeomHard DiskRead Error'
 INVALID_DD = b'Nothing to see here!'
@@ -47,16 +47,14 @@ def __call__(self, args, encoding=None):
 
         elif self.args == ['lsblk', '-spnlo', 'name', BOOT_PARTITION]:
             stdout = BOOT_DEVICE
-        elif self.args == ['lsblk', '-spnlo', 'name', '/dev/sda1']:
-            stdout = '/dev/sda'
-        elif self.args == ['lsblk', '-spnlo', 'name', '/dev/sdb1']:
-            stdout = '/dev/sdb'
+        elif self.args[:-1] == ['lsblk', '-spnlo', 'name']:
+            stdout = self.args[-1][:-1]
 
         return {'stdout': stdout}
 
 
 def open_mocked(fn, flags):
-    if fn == BOOT_DEVICE or fn in MD_BOOT_REAL_DEVICES:
+    if fn == BOOT_DEVICE or fn in MD_BOOT_DEVICES_WITH_GRUB:
         path = os.path.join(CUR_DIR, 'grub_valid')
     else:
         path = os.path.join(CUR_DIR, 'grub_invalid')
@@ -154,7 +152,16 @@ def test_get_grub_devices_one_device(monkeypatch):
     assert 'GRUB is installed on {}'.format(",".join(result)) in api.current_logger.infomsg
 
 
-def test_get_grub_devices_raid_device(monkeypatch):
+@pytest.mark.parametrize(
+    ',component_devs,expected',
+    [
+        (['/dev/sda1', '/dev/sdb1'], MD_BOOT_DEVICES_WITH_GRUB),
+        (['/dev/sda1', '/dev/sdb1', '/dev/sdc1', '/dev/sdd1'], MD_BOOT_DEVICES_WITH_GRUB),
+        (['/dev/sda2', '/dev/sdc1'], ['/dev/sda']),
+        (['/dev/sdd3', '/dev/sdb2'], ['/dev/sdb']),
+    ]
+)
+def test_get_grub_devices_raid_device(monkeypatch, component_devs, expected):
     run_mocked = RunMocked(boot_on_raid=True)
     monkeypatch.setattr(grub, 'run', run_mocked)
     monkeypatch.setattr(os, 'open', open_mocked)
@@ -165,12 +172,12 @@ def test_get_grub_devices_raid_device(monkeypatch):
 
     def get_component_devices_mocked(raid_dev):
         assert raid_dev == MD_BOOT_DEVICE
-        return ['/dev/sda1', '/dev/sdb1']
+        return component_devs
 
     monkeypatch.setattr(mdraid, 'get_component_devices', get_component_devices_mocked)
 
     result = grub.get_grub_devices()
-    assert grub.run.called == 3  # grub2-probe + 2x lsblk
-    assert ['/dev/sda', '/dev/sdb'] == result
+    assert grub.run.called == 1 + len(component_devs)  # grub2-probe + Nx lsblk
+    assert sorted(expected) == sorted(result)
     assert not api.current_logger.warnmsg
     assert 'GRUB is installed on {}'.format(",".join(result)) in api.current_logger.infomsg