Only verify extensions added by the tests for ETP

dcr/scenario_utils/extensions/CustomScriptExtension.py

@@ -5,13 +5,15 @@
 
 
 class CustomScriptExtension(BaseExtensionTestClass):
+    META_DATA = ExtensionMetaData(
+        publisher='Microsoft.Azure.Extensions',
+        ext_type='CustomScript',
+        version="2.1"
+    )
+
     def __init__(self, extension_name: str):
-        extension_data = ExtensionMetaData(
-            publisher='Microsoft.Azure.Extensions',
-            ext_type='CustomScript',
-            version="2.1",
-            ext_name=extension_name
-        )
+        extension_data = self.META_DATA
+        extension_data.name = extension_name
         super().__init__(extension_data)
 
 

dcr/scenario_utils/extensions/VMAccessExtension.py

@@ -8,14 +8,15 @@
 
 
 class VMAccessExtension(BaseExtensionTestClass):
+    META_DATA = ExtensionMetaData(
+        publisher='Microsoft.OSTCExtensions',
+        ext_type='VMAccessForLinux',
+        version="1.5"
+    )
 
     def __init__(self, extension_name: str):
-        extension_data = ExtensionMetaData(
-            publisher='Microsoft.OSTCExtensions',
-            ext_type='VMAccessForLinux',
-            version="1.5",
-            ext_name=extension_name
-        )
+        extension_data = self.META_DATA
+        extension_data.name = extension_name
         super().__init__(extension_data)
         self.public_key, self.private_key_file = generate_ssh_key_pair('dcr_py')
         self.user_name = f'dcr{random_alphanum(length=8)}'

dcr/scenario_utils/models.py

@@ -11,7 +11,7 @@ class VMModelType(Enum):
 
 
 class ExtensionMetaData:
-    def __init__(self, publisher: str, ext_type: str, version: str, ext_name: str):
+    def __init__(self, publisher: str, ext_type: str, version: str, ext_name: str = ""):
         self.__publisher = publisher
         self.__ext_type = ext_type
         self.__version = version
@@ -30,9 +30,17 @@ def version(self) -> str:
         return self.__version
 
     @property
-    def name(self) -> str:
+    def name(self):
         return self.__ext_name
 
+    @name.setter
+    def name(self, ext_name):
+        self.__ext_name = ext_name
+
+    @property
+    def handler_name(self):
+        return f"{self.publisher}.{self.ext_type}"
+
 
 class VMMetaData:
 

dcr/scenarios/extension-telemetry-pipeline/etp_helpers.py

@@ -37,7 +37,7 @@ def wait_for_extension_events_dir_empty(timeout=timedelta(minutes=2)):
     raise AssertionError("Extension events dir not empty!")
 
 
-def add_extension_events_and_get_count(bad_event_count=0, no_of_events_per_extension=50):
+def add_extension_events_and_get_count(bad_event_count=0, no_of_events_per_extension=50, extension_names=None):
     print("Creating random extension events now. No of Good Events: {0}, No of Bad Events: {1}".format(
         no_of_events_per_extension - bad_event_count, bad_event_count))
 
@@ -128,7 +128,8 @@ def oversize_file_limit(make_bad_event):
 
     for ext_dir in os.listdir(ext_log_dir):
         events_dir = os.path.join(ext_log_dir, ext_dir, "events")
-        if not os.path.isdir(events_dir):
+        # If specific extensions are provided, only add the events for them
+        if not os.path.isdir(events_dir) or (extension_names is not None and ext_dir not in extension_names):
             continue
 
         new_opr_id = str(uuid.uuid4())

dcr/scenarios/extension-telemetry-pipeline/run.py

@@ -5,16 +5,19 @@
 
 from dcr.scenario_utils.agent_log_parser import parse_agent_log_file
 from dcr.scenario_utils.check_waagent_log import is_data_in_waagent_log, check_waagent_log_for_errors
+from dcr.scenario_utils.extensions.CustomScriptExtension import CustomScriptExtension
+from dcr.scenario_utils.extensions.VMAccessExtension import VMAccessExtension
 from dcr.scenario_utils.test_orchestrator import TestFuncObj
 from dcr.scenario_utils.test_orchestrator import TestOrchestrator
 from etp_helpers import add_extension_events_and_get_count, wait_for_extension_events_dir_empty, \
     get_collect_telemetry_thread_name
 
 
-def add_good_extension_events_and_verify():
+def add_good_extension_events_and_verify(extension_names):
     max_events = random.randint(10, 50)
     print("Creating a total of {0} events".format(max_events))
-    ext_event_count = add_extension_events_and_get_count(no_of_events_per_extension=max_events)
+    ext_event_count = add_extension_events_and_get_count(no_of_events_per_extension=max_events,
+                                                         extension_names=extension_names)
 
     # Ensure that the event collector ran after adding the events
     wait_for_extension_events_dir_empty()
@@ -36,11 +39,12 @@ def add_good_extension_events_and_verify():
         is_data_in_waagent_log("Collected {0} events for extension: {1}".format(good_count, ext_name))
 
 
-def add_bad_events_and_verify_count():
+def add_bad_events_and_verify_count(extension_names):
     max_events = random.randint(15, 50)
     print("Creating a total of {0} events".format(max_events))
     extension_event_count = add_extension_events_and_get_count(bad_event_count=random.randint(5, max_events - 5),
-                                                               no_of_events_per_extension=max_events)
+                                                               no_of_events_per_extension=max_events,
+                                                               extension_names=extension_names)
 
     # Ensure that the event collector ran after adding the events
     wait_for_extension_events_dir_empty()
@@ -62,13 +66,7 @@ def verify_etp_enabled():
     event_dirs = glob.glob(os.path.join("/var/log/azure/", "*", "events"))
     assert event_dirs, "No extension event directories exist!"
 
-    verified = True
-    for event_dir in event_dirs:
-        exists = os.path.exists(event_dir)
-        print("Dir: {0} exists: {1}".format(event_dir, exists))
-        verified = verified and exists
-
-    if not verified:
+    if not all(os.path.exists(event_dir) for event_dir in event_dirs):
         raise AssertionError("Event directory not found for all extensions!")
 
 
@@ -86,10 +84,14 @@ def check_agent_log():
 
 
 if __name__ == '__main__':
+
+    extensions_to_verify = [CustomScriptExtension.META_DATA.handler_name, VMAccessExtension.META_DATA.handler_name]
     tests = [
         TestFuncObj("Verify ETP enabled", verify_etp_enabled, raise_on_error=True, retry=3),
-        TestFuncObj("Add Good extension events and verify", add_good_extension_events_and_verify),
-        TestFuncObj("Add Bad extension events and verify", add_bad_events_and_verify_count),
+        TestFuncObj("Add Good extension events and verify",
+                    lambda: add_good_extension_events_and_verify(extensions_to_verify)),
+        TestFuncObj("Add Bad extension events and verify",
+                    lambda: add_bad_events_and_verify_count(extensions_to_verify)),
         TestFuncObj("Verify all events processed", wait_for_extension_events_dir_empty),
         TestFuncObj("Check Agent log", check_agent_log),
     ]