2 files changed, 55 insertions, 0 deletions
diff --git a/tools/testing/selftests/net/lib/py/ksft.py b/tools/testing/selftests/net/lib/py/ksft.py
index 789433262dc7..3aaa2748a58e 100644
--- a/tools/testing/selftests/net/lib/py/ksft.py
+++ b/tools/testing/selftests/net/lib/py/ksft.py
@@ -6,6 +6,7 @@ import sys
 import time
 import traceback
 from .consts import KSFT_MAIN_NAME
+from .utils import global_defer_queue
 
 KSFT_RESULT = None
 KSFT_RESULT_ALL = True
@@ -108,6 +109,24 @@ def ktap_result(ok, cnt=1, case="", comment=""):
     print(res)
 
 
+def ksft_flush_defer():
+    global KSFT_RESULT
+
+    i = 0
+    qlen_start = len(global_defer_queue)
+    while global_defer_queue:
+        i += 1
+        entry = global_defer_queue.pop()
+        try:
+            entry.exec_only()
+        except:
+            ksft_pr(f"Exception while handling defer / cleanup (callback {i} of {qlen_start})!")
+            tb = traceback.format_exc()
+            for line in tb.strip().split('\n'):
+                ksft_pr("Defer Exception|", line)
+            KSFT_RESULT = False
+
+
 def ksft_run(cases=None, globs=None, case_pfx=None, args=()):
     cases = cases or []
 
@@ -148,6 +167,8 @@ def ksft_run(cases=None, globs=None, case_pfx=None, args=()):
             KSFT_RESULT = False
             cnt_key = 'fail'
 
+        ksft_flush_defer()
+
         if not cnt_key:
             cnt_key = 'pass' if KSFT_RESULT else 'fail'
 
diff --git a/tools/testing/selftests/net/lib/py/utils.py b/tools/testing/selftests/net/lib/py/utils.py
index 405aa510aaf2..72590c3f90f1 100644
--- a/tools/testing/selftests/net/lib/py/utils.py
+++ b/tools/testing/selftests/net/lib/py/utils.py
@@ -66,6 +66,40 @@ class bkg(cmd):
         return self.process(terminate=self.terminate, fail=self.check_fail)
 
 
+global_defer_queue = []
+
+
+class defer:
+    def __init__(self, func, *args, **kwargs):
+        global global_defer_queue
+
+        if not callable(func):
+            raise Exception("defer created with un-callable object, did you call the function instead of passing its name?")
+
+        self.func = func
+        self.args = args
+        self.kwargs = kwargs
+
+        self._queue =  global_defer_queue
+        self._queue.append(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, ex_type, ex_value, ex_tb):
+        return self.exec()
+
+    def exec_only(self):
+        self.func(*self.args, **self.kwargs)
+
+    def cancel(self):
+        self._queue.remove(self)
+
+    def exec(self):
+        self.cancel()
+        self.exec_only()
+
+
 def tool(name, args, json=None, ns=None, host=None):
     cmd_str = name + ' '
     if json: