From a6b018e3864ff313fed36c3804394e2c92ca87b3 Mon Sep 17 00:00:00 2001
From: Garfield Lee Freeman <gfreeman@paloaltonetworks.com>
Date: Wed, 8 Nov 2023 09:43:11 +0100
Subject: [PATCH] feat(PanDevice): add `is_ready()` (#532)

* feat(PanDevice): add `is_ready()`
---
 Makefile           |  2 +-
 panos/base.py      | 33 +++++++++++++++++++++++++++++++++
 tests/test_base.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index e6aed882..72a264b1 100644
--- a/Makefile
+++ b/Makefile
@@ -53,7 +53,7 @@ format:
 
 check-format:
 	isort --recursive --atomic --check-only panos
-	black --check .
+	black --check --diff .
 
 test:
 	pytest
diff --git a/panos/base.py b/panos/base.py
index 1a9cc90c..120b91a8 100644
--- a/panos/base.py
+++ b/panos/base.py
@@ -5589,3 +5589,36 @@ def whoami(self):
 
                 if name is not None and is_self:
                     return name
+
+    def is_ready(self, minutes=None, seconds=None):
+        """Runs "show chassis-ready" until the PAN-OS management plane is up.
+
+        Args:
+            minutes (int): The number of minutes to wait before giving up.
+            seconds (int): The number of seconds to wait before giving up.
+
+        Returns:
+            True if PAN-OS is ready, or False if a timeout was reached.
+        """
+        end = None
+        if minutes is not None or seconds is not None:
+            end = datetime.datetime.now() + datetime.timedelta(
+                minutes=minutes or 0, seconds=seconds or 0,
+            )
+
+        while True:
+            response = None
+            try:
+                response = self.op("show chassis-ready")
+            except (err.PanURLError, pan.xapi.PanXapiError, err.PanDeviceXapiError):
+                pass
+            else:
+                ready_status = response.find(".//result").text.strip()
+                if ready_status.lower() == "yes":
+                    return True
+
+            if end is not None and datetime.datetime.now() >= end:
+                return False
+
+            # Device isn't up yet, retry after sleeping.
+            time.sleep(2)
diff --git a/tests/test_base.py b/tests/test_base.py
index a432ccc2..235a55d4 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -1620,5 +1620,35 @@ def test_delete_extreme_members(self):
         self.assertTrue(dev.xapi.delete.call_count > 2)
 
 
+class TestIsReady(unittest.TestCase):
+    @mock.patch("time.sleep")
+    def test_ok(self, mocksleep):
+        fw = Base.PanDevice("127.0.0.1", "admin", "secret", api_key="apikey")
+        fw.xapi.op = mock.Mock(
+            side_effect=[
+                Err.PanURLError,
+                pan.xapi.PanXapiError,
+                Err.PanXapiError,
+                ET.fromstring("<response><result>yes</result></response>"),
+                ValueError,
+            ],
+        )
+
+        ans = fw.is_ready()
+
+        assert ans == True
+        assert mocksleep.call_count == 3
+
+    @mock.patch("time.sleep")
+    def test_times_out(self, mocksleep):
+        fw = Base.PanDevice("127.0.0.1", "admin", "secret", api_key="apikey")
+        fw.xapi.op = mock.Mock(side_effect=[Err.PanURLError, ValueError,],)
+
+        ans = fw.is_ready(seconds=0)
+
+        assert ans == False
+        assert mocksleep.call_count == 0
+
+
 if __name__ == "__main__":
     unittest.main()