diff --git a/dcr/scenario_utils/extensions/CustomScriptExtension.py b/dcr/scenario_utils/extensions/CustomScriptExtension.py index 32b285a8ce..29df351134 100644 --- a/dcr/scenario_utils/extensions/CustomScriptExtension.py +++ b/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) diff --git a/dcr/scenario_utils/extensions/VMAccessExtension.py b/dcr/scenario_utils/extensions/VMAccessExtension.py index 0cfa3691ad..c84ae12053 100644 --- a/dcr/scenario_utils/extensions/VMAccessExtension.py +++ b/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)}' diff --git a/dcr/scenario_utils/models.py b/dcr/scenario_utils/models.py index d8825e4acf..806c830c12 100644 --- a/dcr/scenario_utils/models.py +++ b/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: diff --git a/dcr/scenarios/extension-telemetry-pipeline/etp_helpers.py b/dcr/scenarios/extension-telemetry-pipeline/etp_helpers.py index 1e5ad6e1dc..9fd06587ea 100644 --- a/dcr/scenarios/extension-telemetry-pipeline/etp_helpers.py +++ b/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()) diff --git a/dcr/scenarios/extension-telemetry-pipeline/run.py b/dcr/scenarios/extension-telemetry-pipeline/run.py index 9ae7d63b46..3bff11e111 100644 --- a/dcr/scenarios/extension-telemetry-pipeline/run.py +++ b/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), ]