From 211859b11ffacd2d1d27e025eeb91433e15d6fb2 Mon Sep 17 00:00:00 2001
From: Andrew Leech <andrew.leech@planetinnovation.com.au>
Date: Tue, 28 Jun 2022 10:16:48 +1000
Subject: [PATCH] tests/run-multitests.py: Add ability to test instance over
 reboot.

The device-under-test should use `multitest.expect_reboot()` to indicate
that it will reboot.

Signed-off-by: Andrew Leech <andrew.leech@planetinnovation.com.au>
---
 tests/run-multitests.py | 50 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/tests/run-multitests.py b/tests/run-multitests.py
index 2fd5f027e7..b316ce69f6 100755
--- a/tests/run-multitests.py
+++ b/tests/run-multitests.py
@@ -90,6 +90,9 @@ class multitest:
             except:
                 ip = HOST_IP
         return ip
+    @staticmethod
+    def expect_reboot(resume, delay_ms=0):
+        print("WAIT_FOR_REBOOT", resume, delay_ms)
 
 {}
 
@@ -378,6 +381,21 @@ def run_test_on_instances(test_file, num_instances, instances):
                 last_read_time[idx] = time.time()
                 if out is not None and not any(m in out for m in IGNORE_OUTPUT_MATCHES):
                     trace_instance_output(idx, out)
+                    if out.startswith("WAIT_FOR_REBOOT"):
+                        _, resume, delay_ms = out.split(" ")
+
+                        if wait_for_reboot(instance, delay_ms):
+                            # Restart the test code, resuming from requested instance block
+                            if not resume.startswith("instance{}".format(idx)):
+                                raise SystemExit(
+                                    'ERROR: resume function must start with "instance{}"'.format(
+                                        idx
+                                    )
+                                )
+                            append_code = APPEND_CODE_TEMPLATE.format(injected_globals, resume[8:])
+                            instance.start_file(test_file, append=append_code)
+                            last_read_time[idx] = time.time()
+
                     if out.startswith("BROADCAST "):
                         for instance2 in instances:
                             if instance2 is not instance:
@@ -406,6 +424,38 @@ def run_test_on_instances(test_file, num_instances, instances):
     return error, skip, output_str
 
 
+def wait_for_reboot(instance, extra_timeout_ms=0):
+    # Monitor device responses for reboot banner, waiting for idle.
+    extra_timeout = float(extra_timeout_ms) * 1000
+    INITIAL_TIMEOUT = 1 + extra_timeout
+    FULL_TIMEOUT = 5 + extra_timeout
+    t_start = t_last_activity = time.monotonic()
+    while True:
+        t = time.monotonic()
+        out, err = instance.readline()
+        if err is not None:
+            print("Reboot: communication error", err)
+            return False
+        if out:
+            t_last_activity = t
+            # Check for reboot banner, see py/pyexec.c "reset friendly REPL"
+            if re.match(r"^MicroPython v\d+\.\d+\.\d+.* on .*; .* with .*$", out):
+                time.sleep(0.1)
+                break
+
+        if t_last_activity == t_start:
+            if t - t_start > INITIAL_TIMEOUT:
+                print("Reboot: missed initial Timeout")
+                return False
+        else:
+            if t - t_start > FULL_TIMEOUT:
+                print("Reboot: Timeout")
+                return False
+
+    instance.pyb.enter_raw_repl()
+    return True
+
+
 def print_diff(a, b):
     a_fd, a_path = tempfile.mkstemp(text=True)
     b_fd, b_path = tempfile.mkstemp(text=True)