[action] [PR:19246] Updating test_ecn_config_update.py to fix nightly test failures

tests/generic_config_updater/test_ecn_config_update.py

@@ -1,4 +1,5 @@
 import ast
+from functools import cmp_to_key
 import logging
 import pytest
 
@@ -45,45 +46,99 @@ def ensure_dut_readiness(duthost):
         delete_checkpoint(duthost)
 
 
-def ensure_application_of_updated_config(duthost, configdb_field, values):
+def get_asic_db_values(duthost, fields):
+    """
+    Args:
+        duthost: DUT host object
+        fields: CONFIG DB field(s) under test
+
+    Returns:
+        A dictionary where keys are WRED profile OIDs in ASIC DB and values are the field-value pairs
+        for the fields in configdb_field.
+    """
+    wred_objects = duthost.shell('sonic-db-cli ASIC_DB keys *WRED*')["stdout"]
+    wred_objects = wred_objects.split("\n")
+    asic_db_values = {}
+    for wred_object in wred_objects:
+        oid = wred_object[wred_object.rfind(':') + 1:]
+        asic_db_values[oid] = {}
+        wred_data = duthost.shell('sonic-db-cli ASIC_DB hgetall {}'.format(wred_object))["stdout"]
+        if "NULL" in wred_data:
+            continue
+        wred_data = ast.literal_eval(wred_data)
+        for field in fields:
+            value = int(wred_data[WRED_MAPPING[field]])
+            asic_db_values[oid][field] = value
+    return asic_db_values
+
+
+def get_wred_objects(duthost):
+    """
+    Args:
+        duthost: DUT host object
+
+    Returns:
+        A list of WRED profile objects in ASIC DB.
+    """
+    wred_objects = duthost.shell('sonic-db-cli ASIC_DB keys *WRED*')["stdout"]
+    wred_objects = wred_objects.split("\n")
+    return wred_objects
+
+
+def dict_compare(fields):
+    """
+    Compares two dictionaries for equality based on a subset of keys.
+
+    Args:
+        fields: The keys to compare.
+
+    Returns:
+        A function that compares two dictionaries.
+    """
+    def compare(dict1, dict2):
+        for field in fields:
+            if dict1.get(field, 0) < dict2.get(field, 0):
+                return -1
+            elif dict1.get(field, 0) > dict2.get(field, 0):
+                return 1
+        # If all compared fields are equal, return 0
+        return 0
+
+    return compare
+
+
+def ensure_application_of_updated_config(duthost, fields, new_values):
     """
     Ensures application of the JSON patch config update
 
     Args:
         duthost: DUT host object
-        configdb_field: config db field(s) under test
-        values: expected value(s) of configdb_field
+        fields: config db field(s) under test
+        new_values: expected value(s) of fields. It is a dictionary where keys are WRED profile names
+                    and values are dictionaries of field-value pairs for all fields in fields.
     """
-    def _confirm_value_in_asic_db():
-        wred_objects = duthost.shell('sonic-db-cli ASIC_DB keys *WRED*')["stdout"]
-        wred_objects = wred_objects.split("\n")
-        if (len(wred_objects) > 1):
-            for wred_object in wred_objects:
-                wred_data = duthost.shell('sonic-db-cli ASIC_DB hgetall {}'.format(wred_object))["stdout"]
-                if ('NULL' in wred_data):
-                    continue
-                wred_data = ast.literal_eval(wred_data)
-                for field, value in zip(configdb_field.split(','), values.split(',')):
-                    if value != wred_data[WRED_MAPPING[field]]:
-                        return False
-                return True
-            return False
-        else:
-            wred_data = duthost.shell('sonic-db-cli ASIC_DB hgetall {}'.format(wred_objects[0]))["stdout"]
-            wred_data = ast.literal_eval(wred_data)
-            for field, value in zip(configdb_field.split(','), values.split(',')):
-                if value != wred_data[WRED_MAPPING[field]]:
-                    return False
-            return True
-
-    logger.info("Validating fields in ASIC DB...")
+    # Since there is no direct way to obtain the WRED profile name to oid mapping, we will just make sure
+    # that the set of values in ASIC DB matches the set of values in CONFIG DB.
+    def validate_wred_objects_in_asic_db():
+        asic_db_values = get_asic_db_values(duthost, fields)
+        asic_db_values_list = sorted(list(asic_db_values.values()), key=cmp_to_key(dict_compare(fields)))
+        new_values_list = sorted(list(new_values.values()), key=cmp_to_key(dict_compare(fields)))
+        return asic_db_values_list == new_values_list
+
+    logger.info("Validating WRED objects in ASIC DB...")
     pytest_assert(
-        wait_until(READ_ASICDB_TIMEOUT, READ_ASICDB_INTERVAL, 0, _confirm_value_in_asic_db),
+        wait_until(READ_ASICDB_TIMEOUT, READ_ASICDB_INTERVAL, 0, validate_wred_objects_in_asic_db),
         "ASIC DB does not properly reflect newly configured field(s): {} expected value(s): {}"
-        .format(configdb_field, values)
+        .format(fields, new_values)
     )
 
 
+def get_wred_profiles(duthost):
+    wred_profiles = duthost.shell("sonic-db-cli CONFIG_DB keys 'WRED_PROFILE|*' | cut -d '|' -f 2")["stdout"]
+    wred_profiles = wred_profiles.split('\n')
+    return wred_profiles
+
+
 @pytest.mark.parametrize("configdb_field", ["green_min_threshold", "green_max_threshold", "green_drop_probability",
                          "green_min_threshold,green_max_threshold,green_drop_probability"])
 @pytest.mark.parametrize("operation", ["replace"])
@@ -92,28 +147,37 @@ def test_ecn_config_updates(duthost, ensure_dut_readiness, configdb_field, opera
     logger.info("tmpfile {} created for json patch of field: {} and operation: {}"
                 .format(tmpfile, configdb_field, operation))
 
+    fields = configdb_field.split(',')
+    wred_profiles = get_wred_profiles(duthost)
+    if not wred_profiles:
+        pytest.skip("No WRED profiles found in CONFIG_DB, skipping test.")
     json_patch = list()
-    values = list()
-    ecn_data = duthost.shell('sonic-db-cli CONFIG_DB hgetall "WRED_PROFILE|AZURE_LOSSLESS"')['stdout']
-    ecn_data = ast.literal_eval(ecn_data)
-    for field in configdb_field.split(','):
-        value = int(ecn_data[field]) + 1
-        values.append(str(value))
-
-        logger.info("value to be added to json patch: {}, operation: {}, field: {}"
-                    .format(value, operation, field))
-
-        json_patch.append(
-                          {"op": "{}".format(operation),
-                           "path": "/WRED_PROFILE/AZURE_LOSSLESS/{}".format(field),
-                           "value": "{}".format(value)})
+    # new_values is a dictionary from WRED profile name to its field-value mapping (with new values)
+    # for the fields in configdb_field.
+    new_values = {}
+    # Creating a JSON patch for all WRED profiles in CONFIG_DB.
+    for wred_profile in wred_profiles:
+        ecn_data = duthost.shell(f"sonic-db-cli CONFIG_DB hgetall 'WRED_PROFILE|{wred_profile}'")["stdout"]
+        ecn_data = ast.literal_eval(ecn_data)
+        new_values[wred_profile] = {}
+        for field in fields:
+            value = int(ecn_data[field]) + 1
+            new_values[wred_profile][field] = value
+
+            logger.info("value to be added to json patch: {}, operation: {}, field: {}"
+                        .format(value, operation, field))
+
+            json_patch.append(
+                              {"op": "{}".format(operation),
+                               "path": f"/WRED_PROFILE/{wred_profile}/{field}",
+                               "value": "{}".format(value)})
 
     json_patch = format_json_patch_for_multiasic(duthost=duthost, json_data=json_patch)
     try:
         output = apply_patch(duthost, json_data=json_patch, dest_file=tmpfile)
         if is_valid_platform_and_version(duthost, "WRED_PROFILE", "ECN tuning", operation):
             expect_op_success(duthost, output)
-            ensure_application_of_updated_config(duthost, configdb_field, ",".join(values))
+            ensure_application_of_updated_config(duthost, fields, new_values)
         else:
             expect_op_failure(output)
     finally: