From dec26dfc63429044de0bb6304e20e3fbbd244821 Mon Sep 17 00:00:00 2001
From: Anton Smorodskyi <asmorodskyi@suse.com>
Date: Fri, 24 May 2024 11:31:24 +0200
Subject: [PATCH] Add collecting stat for EC2

---
 ocw/lib/dump_state.py | 27 +++++++++++++++++++++++++++
 ocw/lib/ec2.py        | 21 +++++++++++++++++++++
 ocw/lib/influx.py     |  1 +
 tests/test_ec2.py     | 31 ++++++++++++++++++++++---------
 4 files changed, 71 insertions(+), 9 deletions(-)

diff --git a/ocw/lib/dump_state.py b/ocw/lib/dump_state.py
index 27b3a5bf..f7256626 100644
--- a/ocw/lib/dump_state.py
+++ b/ocw/lib/dump_state.py
@@ -1,7 +1,9 @@
+import os
 import logging
 import traceback
 from webui.PCWConfig import PCWConfig
 from ocw.lib.azure import Azure
+from ocw.lib.ec2 import EC2
 from ocw.enums import ProviderChoice
 from ocw.lib.influx import Influx
 
@@ -9,6 +11,12 @@
 
 
 def dump_state():
+    if os.getenv("INFLUX_TOKEN") is None:
+        logger.warning("INFLUX_TOKEN is not set, dumping state is not possible")
+        return
+    if not PCWConfig.has("influxdb/url"):
+        logger.warning("pcw.ini missing influxdb configuration, dumping state is not possible")
+        return
     for namespace in PCWConfig.get_namespaces_for("influxdb"):
         try:
             providers = PCWConfig.get_providers_for("influxdb", namespace)
@@ -38,6 +46,25 @@ def dump_state():
                     namespace,
                     Azure(namespace).get_img_versions_count,
                 )
+            if ProviderChoice.EC2 in providers:
+                Influx().dump_resource(
+                    ProviderChoice.EC2.value,
+                    Influx.VMS_QUANTITY,
+                    namespace,
+                    EC2(namespace).count_all_instances
+                )
+                Influx().dump_resource(
+                    ProviderChoice.EC2.value,
+                    Influx.IMAGES_QUANTITY,
+                    namespace,
+                    EC2(namespace).count_all_images
+                )
+                Influx().dump_resource(
+                    ProviderChoice.EC2.value,
+                    Influx.VOLUMES_QUANTITY,
+                    namespace,
+                    EC2(namespace).count_all_volumes
+                )
         except Exception:
             logger.exception(
                 "[%s] Dump state failed!: \n %s", namespace, traceback.format_exc()
diff --git a/ocw/lib/ec2.py b/ocw/lib/ec2.py
index dab2b383..7a43a0c7 100644
--- a/ocw/lib/ec2.py
+++ b/ocw/lib/ec2.py
@@ -116,6 +116,13 @@ def volume_protected(self, volume: dict) -> bool:
     def list_instances(self, region: str) -> list:
         return list(self.ec2_resource(region).instances.all())
 
+    def count_all_instances(self) -> int:
+        instance_quantity = 0
+        for region in self.all_regions:
+            instances = self.list_instances(region=region)
+            instance_quantity += len(instances)
+        return instance_quantity
+
     def get_all_regions(self) -> list:
         regions_resp = self.ec2_client(EC2.default_region).describe_regions()
         regions = [region['RegionName'] for region in regions_resp['Regions']]
@@ -324,6 +331,20 @@ def report_cleanup_results(self, vpc_errors: list, vpc_notify: list, vpc_locked:
         if len(vpc_locked) > 0:
             send_mail('VPC deletion locked by running VMs', '\n'.join(vpc_locked))
 
+    def count_all_images(self) -> int:
+        all_images_cnt = 0
+        for region in self.all_regions:
+            response = self.ec2_client(region).describe_images(Owners=['self'])
+            all_images_cnt += len(response['Images'])
+        return all_images_cnt
+
+    def count_all_volumes(self) -> int:
+        all_volumes_cnt = 0
+        for region in self.all_regions:
+            response = self.ec2_client(region).describe_volumes()
+            all_volumes_cnt += len(response['Volumes'])
+        return all_volumes_cnt
+
     def cleanup_images(self, valid_period_days: float) -> None:
         self.log_dbg('Call cleanup_images')
         for region in self.all_regions:
diff --git a/ocw/lib/influx.py b/ocw/lib/influx.py
index 774ebb50..9b921c5f 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"
+    VOLUMES_QUANTITY: str = "volumes_quanity"
     IMAGE_VERSION_QUANTITY: str = "img_version_quantity"
     NAMESPACE_TAG: str = "namespace"
 
diff --git a/tests/test_ec2.py b/tests/test_ec2.py
index 34f4fe41..44398d96 100644
--- a/tests/test_ec2.py
+++ b/tests/test_ec2.py
@@ -83,6 +83,22 @@ def mock_local_get_feature_property(feature: str, property: str, namespace: str
     monkeypatch.setattr('smtplib.SMTP', lambda arg1, arg2: MockedSMTP())
     return EC2('fake')
 
+class MockedInstances:
+    is_empty = False
+
+    def __init__(self, subnet=True) -> None:
+        self.subnet = subnet
+
+    def all(self):
+        if self.subnet:
+            if MockedInstances.is_empty:
+                return []
+            else:
+                return ['subnet']
+        else:
+            return 5
+
+
 
 class MockedEC2Client():
     response = {}
@@ -127,6 +143,9 @@ class MockedEC2Client():
             }
         ]}
 
+    def __init__(self, region):
+        self.instances = MockedInstances(subnet=False)
+
     ec2_snapshots = {snapshotid_to_delete: 'snapshot', snapshotid_i_have_ami: 'snapshot'}
 
     def describe_images(self, *args, **kwargs):
@@ -185,15 +204,6 @@ def sendmail(self, sender_email, receiver_email, mimetext):
         MockedSMTP.mimetext = mimetext
 
 
-class MockedInstances:
-    is_empty = False
-
-    def all(self):
-        if MockedInstances.is_empty:
-            return []
-        else:
-            return ['subnet']
-
 
 class MockedInterface:
     delete_called = False
@@ -507,3 +517,6 @@ def mocked_get_boolean(config_path, field=None):
     ec2_patch.cleanup_all()
 
     assert called_stack == ['cleanup_images', 'cleanup_snapshots', 'cleanup_volumes', 'cleanup_vpcs']
+
+def test_list_all_instances(ec2_patch):
+    assert ec2_patch.count_all_instances() == 5