Merge pull request #146 from sai1274/master

Added Master Arbitration support for Set
akarneliuk · Mar 1, 2024 · 038564f · 038564f
2 parents b540b97 + 612cb65
commit 038564f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/pygnmi/client.py b/pygnmi/client.py
@@ -553,6 +553,7 @@ def set(
         encoding: str = None,
         prefix: str = "",
         target: str = None,
+        extension: dict = None,
     ):
         """
         Changing the configuration on the destination network elements.
@@ -583,6 +584,7 @@ def set(
 
         # Set the encoding to auto-discovered value, unless overridden
         encoding = encoding or self.__encoding
+        gnmi_extension = get_gnmi_extension(ext=extension)
 
         # Gnmi PREFIX
         try:
@@ -630,9 +632,12 @@ def set(
                                            encoding=encoding,
                                            datatype='config')
 
-            gnmi_message_request = SetRequest(
-                prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg
-            )
+            if gnmi_extension:
+                gnmi_message_request = SetRequest(prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg, extension=[gnmi_extension])
+            else:
+                gnmi_message_request = SetRequest(
+                    prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg
+                )
             debug_gnmi_msg(self.__debug, gnmi_message_request, "gNMI request")
 
             gnmi_message_response = self.__stub.Set(gnmi_message_request, metadata=self.__metadata)

diff --git a/pygnmi/create_gnmi_extension.py b/pygnmi/create_gnmi_extension.py
@@ -36,6 +36,14 @@ def get_gnmi_extension(ext: dict = None) -> list:
                 result.history.range.start = _get_time_ns_epoch(ext["history"]["range"]["start"])
                 result.history.range.end = _get_time_ns_epoch(ext["history"]["range"]["end"])
 
+    if "master_arbitration" in ext:
+
+        result.master_arbitration.election_id.high = ext["master_arbitration"]["election_id"]["high"]
+        result.master_arbitration.election_id.low = ext["master_arbitration"]["election_id"]["low"]
+
+        if "role" in ext["master_arbitration"]:
+            result.master_arbitration.role.id = ext["master_arbitration"]["role"]["id"]
+
     return result