Enumerate VDS queries using a common function

Change-Id: I32a9844465656f77686a2e92822b180aa9f3a7fb

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"

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()
+        for unk in _enumerate(pack.QueryVolumes()):
-        while True:
-            (unk, c) = enum.Next(1)
-            if not c:
-                break
             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()
+        for unk in _enumerate(pack.QueryDisks()):
-        while True:
-            (unk, c) = enum.Next(1)
-            if not c:
-                break
             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 = []
+        return [unk.QueryInterface(vds.IVdsSwProvider)
-        enum = svc.QueryProviders(vds.VDS_QUERY_SOFTWARE_PROVIDERS)
+                for unk in _enumerate(
-        while True:
+                    svc.QueryProviders(vds.VDS_QUERY_SOFTWARE_PROVIDERS))]
-            (unk, c) = enum.Next(1)
-            if not c:
-                break
-            providers.append(unk.QueryInterface(vds.IVdsSwProvider))
-        return providers
 
     def _query_packs(self, provider):
-        packs = []
+        return [unk.QueryInterface(vds.IVdsPack)
-        enum = provider.QueryPacks()
+                for unk in _enumerate(provider.QueryPacks())]
-        while True:
-            (unk, c) = enum.Next(1)
-            if not c:
-                break
-            packs.append(unk.QueryInterface(vds.IVdsPack))
-        return packs
 
     def extend_volumes(self, volume_indexes=None):
         svc = vds.load_vds_service()