From 5b9dde99528a219ca6952822b277de28b1f2fe03 Mon Sep 17 00:00:00 2001 From: Cosmin Poieana Date: Mon, 14 Sep 2015 04:31:02 +0300 Subject: [PATCH] Enumerate VDS queries using a common function Change-Id: I32a9844465656f77686a2e92822b180aa9f3a7fb --- .../storage/test_vds_storage_manager.py | 12 +++++- .../windows/storage/vds_storage_manager.py | 42 +++++++------------ 2 files changed, 26 insertions(+), 28 deletions(-) diff --git a/cloudbaseinit/tests/utils/windows/storage/test_vds_storage_manager.py b/cloudbaseinit/tests/utils/windows/storage/test_vds_storage_manager.py index 442e1412..1dc9d2c6 100644 --- a/cloudbaseinit/tests/utils/windows/storage/test_vds_storage_manager.py +++ b/cloudbaseinit/tests/utils/windows/storage/test_vds_storage_manager.py @@ -38,12 +38,20 @@ class TestVDSStorageManager(unittest.TestCase): self._module_patcher.start() - vds_store = importlib.import_module( + self.vds_store = importlib.import_module( "cloudbaseinit.utils.windows.storage.vds_storage_manager") - self._vds_storage_manager = vds_store.VDSStorageManager() + self._vds_storage_manager = self.vds_store.VDSStorageManager() self.addCleanup(self._module_patcher.stop) + def test_enumerate(self): + query = mock.Mock() + items = [1, 2, 3] + query.Next.side_effect = [(item, True) for item in items] + \ + [(None, None)] + response = list(self.vds_store._enumerate(query)) + self.assertEqual(items, response) + @mock.patch("cloudbaseinit.utils.windows.storage.vds_storage_manager" ".VDSStorageManager._get_volume_index") @mock.patch("cloudbaseinit.utils.windows.storage.vds_storage_manager" diff --git a/cloudbaseinit/utils/windows/storage/vds_storage_manager.py b/cloudbaseinit/utils/windows/storage/vds_storage_manager.py index 29865fd4..594ee863 100644 --- a/cloudbaseinit/utils/windows/storage/vds_storage_manager.py +++ b/cloudbaseinit/utils/windows/storage/vds_storage_manager.py @@ -27,13 +27,18 @@ ole32.CoTaskMemFree.restype = None ole32.CoTaskMemFree.argtypes = [ctypes.c_void_p] +def _enumerate(query): + """Enumerate VDS service queries.""" + while True: + unk, avail = query.Next(1) + if not avail: + return + yield unk + + class VDSStorageManager(base.BaseStorageManager): def _extend_volumes(self, pack, volume_indexes): - enum = pack.QueryVolumes() - while True: - (unk, c) = enum.Next(1) - if not c: - break + for unk in _enumerate(pack.QueryVolumes()): volume = unk.QueryInterface(vds.IVdsVolume) volume_prop = volume.GetProperties() try: @@ -82,11 +87,7 @@ class VDSStorageManager(base.BaseStorageManager): def _get_volume_extents_to_resize(self, pack, volume_id): volume_extents = [] - enum = pack.QueryDisks() - while True: - (unk, c) = enum.Next(1) - if not c: - break + for unk in _enumerate(pack.QueryDisks()): disk = unk.QueryInterface(vds.IVdsDisk) (extents_p, num_extents) = disk.QueryExtents() @@ -118,24 +119,13 @@ class VDSStorageManager(base.BaseStorageManager): return [ve for ve in volume_extents if ve[1] > 0] def _query_providers(self, svc): - providers = [] - enum = svc.QueryProviders(vds.VDS_QUERY_SOFTWARE_PROVIDERS) - while True: - (unk, c) = enum.Next(1) - if not c: - break - providers.append(unk.QueryInterface(vds.IVdsSwProvider)) - return providers + return [unk.QueryInterface(vds.IVdsSwProvider) + for unk in _enumerate( + svc.QueryProviders(vds.VDS_QUERY_SOFTWARE_PROVIDERS))] def _query_packs(self, provider): - packs = [] - enum = provider.QueryPacks() - while True: - (unk, c) = enum.Next(1) - if not c: - break - packs.append(unk.QueryInterface(vds.IVdsPack)) - return packs + return [unk.QueryInterface(vds.IVdsPack) + for unk in _enumerate(provider.QueryPacks())] def extend_volumes(self, volume_indexes=None): svc = vds.load_vds_service()