Nokia srl prompt stripping (#3531)

Co-authored-by: Jeff Kala <[email protected]>
ktbyers · Nov 8, 2024 · 5e77f23 · 5e77f23
1 parent 3330854
commit 5e77f23
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 1 deletion.
diff --git a/netmiko/nokia/nokia_srl.py b/netmiko/nokia/nokia_srl.py
@@ -53,6 +53,32 @@ def session_preparation(self) -> None:
             self.disable_paging(command=command, cmd_verify=True, pattern=r"#")
         self.set_base_prompt()
 
+    def strip_prompt(self, *args: Any, **kwargs: Any) -> str:
+        """Strip the prompt and the additional context line"""
+        a_string = super().strip_prompt(*args, **kwargs)
+        return self._strip_context_items(a_string)
+
+    def _strip_context_items(self, a_string: str) -> str:
+        """Strip NokiaSRL-specific output.
+
+        Nokia will put extra context in the 1st line of the prompt, such as:
+        --{ running }--[  ]--
+        --{ candidate private private-admin }--[  ]--
+        --{ candidate private private-admin }--[  ]--
+
+        This method removes those lines.
+        """
+        strings_to_strip = [
+            r"--{.*\B",
+        ]
+
+        response_list = a_string.split(self.RESPONSE_RETURN)
+        last_line = response_list[-1]
+        for pattern in strings_to_strip:
+            if re.search(pattern, last_line, flags=re.I):
+                return self.RESPONSE_RETURN.join(response_list[:-1])
+        return a_string
+
     def set_base_prompt(
         self,
         pri_prompt_terminator: str = "#",

diff --git a/tests/unit/test_base_connection.py b/tests/unit/test_base_connection.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-
+import pytest
 import time
 from os.path import dirname, join
 from threading import Lock
@@ -553,3 +553,71 @@ def test_disable_sha2_fix():
     assert {"rsa-sha2-512", "rsa-sha2-256"} & allowed_pubkeys == set()
 
     connection.disconnect()
+
+
+TEST_CASES = [
+    ("some important data\n--{ running }--[  ]--\nA:srl1#", "some important data"),
+    (
+        "more data\nsome important data\n--{ running }--[  ]--\nA:srl1#",
+        "more data\nsome important data",
+    ),
+    (
+        "more data\nsome important data\n--{ candidate private private-admin }--[  ]--\nA:srl1#",
+        "more data\nsome important data",
+    ),
+    (
+        "more data\nsome important data\n--{ candidate private private-admin }--[  ]--\nA:srl1#",
+        "more data\nsome important data",
+    ),
+    (
+        """
+{
+  "basic system info": {
+    "Hostname": "srl1",
+    "Chassis Type": "7220 IXR-D2L",
+    "Part Number": "Sim Part No.",
+    "Serial Number": "Sim Serial No.",
+    "System HW MAC Address": "1A:03:00:FF:00:00",
+    "OS": "SR Linux",
+    "Software Version": "v24.7.2",
+    "Build Number": "319-g64b71941f7",
+    "Architecture": "<Unknown>",
+    "Last Booted": "2024-11-01T17:21:00.164Z",
+    "Total Memory": "<Unknown>",
+    "Free Memory": "<Unknown>"
+  }
+}
+
+--{ running }--[  ]--
+A:srl1#""",
+        """
+{
+  "basic system info": {
+    "Hostname": "srl1",
+    "Chassis Type": "7220 IXR-D2L",
+    "Part Number": "Sim Part No.",
+    "Serial Number": "Sim Serial No.",
+    "System HW MAC Address": "1A:03:00:FF:00:00",
+    "OS": "SR Linux",
+    "Software Version": "v24.7.2",
+    "Build Number": "319-g64b71941f7",
+    "Architecture": "<Unknown>",
+    "Last Booted": "2024-11-01T17:21:00.164Z",
+    "Total Memory": "<Unknown>",
+    "Free Memory": "<Unknown>"
+  }
+}
+""",
+    ),
+]
+
+
+@pytest.mark.parametrize("test_string,expected", TEST_CASES)
+def test_nokiasrl_prompt_stripping(test_string, expected):
+    conn = ConnectHandler(
+        host="testhost",
+        device_type="nokia_srl",
+        auto_connect=False,  # No need to connect for the test purposes
+    )
+    result = conn.strip_prompt(a_string=test_string)
+    assert result == expected