diff --git a/ocw/lib/azure.py b/ocw/lib/azure.py index f92f51a2..35376797 100644 --- a/ocw/lib/azure.py +++ b/ocw/lib/azure.py @@ -18,9 +18,9 @@ class Azure(Provider): def __init__(self, namespace: str): super().__init__(namespace) - self.__resource_group = PCWConfig.get_feature_property('cleanup', 'azure-storage-resourcegroup', namespace) + self.__resource_group: str = str(PCWConfig.get_feature_property('cleanup', 'azure-storage-resourcegroup', namespace)) self.check_credentials() - self.__gallery = PCWConfig.get_feature_property('cleanup', 'azure-gallery-name', namespace) + self.__gallery: str = str(PCWConfig.get_feature_property('cleanup', 'azure-gallery-name', namespace)) def __new__(cls, namespace: str) -> 'Azure': if namespace not in Azure.__instances: @@ -206,3 +206,14 @@ def cleanup_gallery_img_versions(self) -> None: self.compute_mgmt_client().gallery_image_versions.begin_delete( self.__resource_group, gallery.name, image.name, version.name ) + + def get_img_versions_count(self) -> int: + self.log_dbg("Call get_img_versions_count") + gallery = self.compute_mgmt_client().galleries.get(self.__resource_group, self.__gallery) + all_img_versions = 0 + for image_definition in self.compute_mgmt_client().gallery_images.list_by_gallery(self.__resource_group, gallery.name): + img_versions = len(list(self.compute_mgmt_client().gallery_image_versions.list_by_gallery_image( + self.__resource_group, gallery.name, image_definition.name))) + self.log_dbg(f"{image_definition.name} has {img_versions} versions") + all_img_versions += img_versions + return all_img_versions diff --git a/ocw/lib/dump_state.py b/ocw/lib/dump_state.py index 362ed002..f9a1e60c 100644 --- a/ocw/lib/dump_state.py +++ b/ocw/lib/dump_state.py @@ -17,5 +17,6 @@ def dump_state(): Influx().dump_resource(ProviderChoice.AZURE.value, Influx.VMS_QUANTITY, Azure(namespace).list_instances) Influx().dump_resource(ProviderChoice.AZURE.value, Influx.IMAGES_QUANTITY, Azure(namespace).list_images) Influx().dump_resource(ProviderChoice.AZURE.value, Influx.DISK_QUANTITY, Azure(namespace).list_disks) + Influx().dump_resource(ProviderChoice.AZURE.value, Influx.IMAGE_VERSION_QUANTITY, Azure(namespace).get_img_versions_count) except Exception: logger.exception("[%s] Dump state failed!: \n %s", namespace, traceback.format_exc()) diff --git a/ocw/lib/influx.py b/ocw/lib/influx.py index 2a75f096..ce654384 100644 --- a/ocw/lib/influx.py +++ b/ocw/lib/influx.py @@ -17,6 +17,7 @@ class Influx: VMS_QUANTITY: str = "vms_quantity" IMAGES_QUANTITY: str = "images_quantity" DISK_QUANTITY: str = "disk_quantity" + IMAGE_VERSION_QUANTITY: str = "img_version_quantity" def __init__(self) -> None: if self.__client is None: @@ -47,6 +48,12 @@ def write(self, measurement: str, field: str, value: int) -> None: logger.warning("Failed to write to influxdb(record=%s): %s", point, exception) def dump_resource(self, provider: str, field: str, dump_method: Callable) -> None: - items_cnt = len(dump_method()) - logger.debug("%d instances found in %s", items_cnt, provider) + return_value = dump_method() + if isinstance(return_value, list): + items_cnt = len(return_value) + elif isinstance(return_value, int): + items_cnt = return_value + else: + raise ValueError(f"{dump_method} returned unsupported type {type(return_value)}") + logger.debug("%s=%d for %s", field, items_cnt, provider) self.write(provider, field, items_cnt) diff --git a/tests/test_azure.py b/tests/test_azure.py index 0007207f..63d3191c 100644 --- a/tests/test_azure.py +++ b/tests/test_azure.py @@ -58,12 +58,29 @@ def mock_compute_mgmt_client(self): def compute_mgmt_client(): pass compute_mgmt_client.images = lambda: None + compute_mgmt_client.galleries = lambda: None + compute_mgmt_client.gallery_images = lambda: None + compute_mgmt_client.gallery_image_versions = lambda: None + compute_mgmt_client.galleries.get = lambda rg, name: FakeGalleryAndImageDefinition() + compute_mgmt_client.gallery_images.list_by_gallery = lambda rg, name: [FakeGalleryAndImageDefinition()] + compute_mgmt_client.gallery_image_versions.list_by_gallery_image = lambda rg, gallery, definitionname: [1, 2, 3, 4, 5] compute_mgmt_client.images.begin_delete = lambda rg, name: deleted_images.append(name) return compute_mgmt_client monkeypatch.setattr(Azure, 'compute_mgmt_client', mock_compute_mgmt_client) +# This class is faking two unrelated entities: +# 1. Gallery returned by compute_mgmt_client.galleries.get +# 2. ImageDefinition returned by compute_mgmt_client.gallery_images.list_by_gallery +# But because all what we need in this case is "something what has property 'name'" +# it is fine to use same class +class FakeGalleryAndImageDefinition: + + def __init__(self) -> None: + self.name = "name" + + class FakeDisk: def __init__(self, managed_by=None): @@ -369,3 +386,7 @@ def compute_mgmt_client(): # when one of VMs throw azure.core.exceptions.ResourceNotFoundError we returning None ret = azure.get_vm_types_in_resource_group('fire!!!') assert ret is None + + +def test_get_img_versions_count(azure_patch, mock_compute_mgmt_client): + assert Azure('fake').get_img_versions_count() == 5