[PATCH net-next 2/2] selftests: net: py: ensure defer() is only used within a test case

[Jakub Kicinski <kuba@kernel.org>]

I wasted a couple of hours recently after accidentally adding
a defer() from within a function which itself was called as
part of defer(). This leads to an infinite loop of defer().
Make sure this cannot happen and raise a helpful exception.

I understand that the pair of _ksft_defer_arm() calls may
not be the most Pythonic way to implement this, but it's
easy enough to understand.

Signed-off-by: Jakub Kicinski <kuba@kernel.org>
---
 tools/testing/selftests/net/lib/py/ksft.py  | 7 +++++++
 tools/testing/selftests/net/lib/py/utils.py | 3 +++
 2 files changed, 10 insertions(+)

diff --git a/tools/testing/selftests/net/lib/py/ksft.py b/tools/testing/selftests/net/lib/py/ksft.py
index 248cd1a723a3..0a96f88bb60a 100644
--- a/tools/testing/selftests/net/lib/py/ksft.py
+++ b/tools/testing/selftests/net/lib/py/ksft.py
@@ -153,6 +153,11 @@ KSFT_DISRUPTIVE = True
     print(res, flush=True)
 
 
+def _ksft_defer_arm(state):
+    """ Allow or disallow the use of defer() """
+    utils.GLOBAL_DEFER_ARMED = state
+
+
 def ksft_flush_defer():
     global KSFT_RESULT
 
@@ -315,6 +320,7 @@ KsftCaseFunction = namedtuple("KsftCaseFunction",
         comment = ""
         cnt_key = ""
 
+        _ksft_defer_arm(True)
         try:
             func(*args)
         except KsftSkipEx as e:
@@ -332,6 +338,7 @@ KsftCaseFunction = namedtuple("KsftCaseFunction",
                 ksft_pr(f"Stopping tests due to {type(e).__name__}.")
             KSFT_RESULT = False
             cnt_key = 'fail'
+        _ksft_defer_arm(False)
 
         try:
             ksft_flush_defer()
diff --git a/tools/testing/selftests/net/lib/py/utils.py b/tools/testing/selftests/net/lib/py/utils.py
index 2dde34560d65..824f039d384c 100644
--- a/tools/testing/selftests/net/lib/py/utils.py
+++ b/tools/testing/selftests/net/lib/py/utils.py
@@ -142,6 +142,7 @@ import time
 
 
 GLOBAL_DEFER_QUEUE = []
+GLOBAL_DEFER_ARMED = False
 
 
 class defer:
@@ -153,6 +154,8 @@ GLOBAL_DEFER_QUEUE = []
         self.args = args
         self.kwargs = kwargs
 
+        if not GLOBAL_DEFER_ARMED:
+            raise Exception("defer queue not armed, did you use defer() outside of a test case?")
         self._queue = GLOBAL_DEFER_QUEUE
         self._queue.append(self)
 
-- 
2.52.0