From 6c9d756f11c8675ec59389ed2031bb09b64b665e Mon Sep 17 00:00:00 2001
From: Peter Barker <pbarker@barker.dropbear.id.au>
Date: Mon, 6 Jul 2020 10:41:16 +1000
Subject: [PATCH] autotest: validate context nesting

---
 Tools/autotest/ardusub.py |  6 ++++--
 Tools/autotest/common.py  | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/Tools/autotest/ardusub.py b/Tools/autotest/ardusub.py
index 9ef1d8e940..aba7f36eaa 100644
--- a/Tools/autotest/ardusub.py
+++ b/Tools/autotest/ardusub.py
@@ -206,7 +206,10 @@ class AutoTestSub(AutoTest):
         self.progress("Mission OK")
 
     def test_gripper_mission(self):
-        self.context_push()
+        with self.Context(self):
+            self.test_gripper_body()
+
+    def test_gripper_body(self):
         ex = None
         try:
             try:
@@ -225,7 +228,6 @@ class AutoTestSub(AutoTest):
         except Exception as e:
             self.progress("Exception caught")
             ex = e
-        self.context_pop()
         if ex is not None:
             raise ex
 
diff --git a/Tools/autotest/common.py b/Tools/autotest/common.py
index c8c0a16e93..6f97de9de2 100644
--- a/Tools/autotest/common.py
+++ b/Tools/autotest/common.py
@@ -2858,6 +2858,15 @@ class AutoTest(ABC):
                                old_value,
                                add_to_context=False)
 
+    class Context(object):
+        def __init__(self, testsuite):
+            self.testsuite = testsuite
+        def __enter__(self):
+            self.testsuite.context_push()
+        def __exit__(self, type, value, traceback):
+            self.testsuite.context_pop()
+            return False # re-raise any exception
+
     def sysid_thismav(self):
         return 1
 
@@ -3842,6 +3851,7 @@ Also, ignores heartbeats not from our target system'''
         self.send_statustext(prettyname)
         self.start_test(prettyname)
         self.set_current_test_name(name)
+        old_contexts_length = len(self.contexts)
         self.context_push()
 
         start_time = time.time()
@@ -3883,6 +3893,11 @@ Also, ignores heartbeats not from our target system'''
             self.progress('Corefiles detected: %s' % str(corefiles))
             passed = False
 
+        if len(self.contexts) != old_contexts_length:
+            self.progress("context count mismatch (want=%u got=%u)" %
+                          (old_contexts_length, len(self.contexts)))
+            passed = False
+
         if passed:
 #            self.remove_bin_logs() # can't do this as one of the binlogs is probably open for writing by the SITL process.  If we force a rotate before running tests then we can do this.
             pass
@@ -5951,6 +5966,7 @@ switch value'''
         except Exception as e:
             ex = e
         self.mavproxy.send("fence disable\n")
+        self.context_pop()
         if ex is not None:
             raise ex