CMLDEV-12: Add posibility to switch serial console used by PyAts conn…

virl2_client/models/cl_pyats.py

@@ -21,6 +21,7 @@
 from __future__ import annotations
 
 import io
+import re
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -120,6 +121,37 @@ def sync_testbed(self, username: str, password: str) -> None:
         self._testbed = self._load_pyats_testbed(testbed_yaml)
         self.set_termserv_credentials(username, password)
 
+    def switch_pyats_serial_console(self, node_label: str, console_number: int) -> None:
+        """
+        Switch to different serial console that is used to execute PyAts commands
+        should be executed after sync_testbed
+        and re-executed after every sync_testbed call.
+
+        :param node_label: The label/title of the device.
+        :param console_number: The serial console number to be used for PyAts.
+        """
+
+        try:
+            pyats_device: Device = self._testbed.devices[node_label]
+        except KeyError:
+            raise PyatsDeviceNotFound(node_label)
+
+        node = self._lab.get_node_by_label(node_label)
+
+        # will raise API error inside if console does not exist for device
+        node.console_key(console_number)
+
+        old_connect_command = pyats_device.connections["a"]["command"]
+
+        pattern = rf"({re.escape(node_label)})/\d+"
+        new_console_cfg = f"/{console_number}"
+
+        new_connect_command = re.sub(
+            pattern, rf"\1{new_console_cfg}", old_connect_command
+        )
+
+        pyats_device.connections["a"]["command"] = new_connect_command
+
     def set_termserv_credentials(
         self,
         username: str | None = None,

virl2_client/models/node.py

@@ -717,14 +717,16 @@ def console_logs(self, console_id: int, lines: int | None = None) -> dict:
         return self._session.get(url).json()
 
     @check_stale
-    def console_key(self) -> str:
+    def console_key(self, console_number: int = 0) -> str:
         """
         Get the console key of the node.
 
+        :param console_number: The console number (defaults to 0).
         :returns: The console key.
         """
+        params = {"line": console_number}
         url = self._url_for("console_key")
-        return self._session.get(url).json()
+        return self._session.get(url, params=params).json()
 
     @check_stale
     def vnc_key(self) -> str: