selftests: net: support use of NetdevSimDev under "with" in python

Using "with" on an entire driver test env is supported already,
but it's also useful to use "with" on an individual nsim.

Reviewed-by: Petr Machata <petrm@nvidia.com>
Link: https://lore.kernel.org/r/20240412141436.828666-6-kuba@kernel.org
Signed-off-by: Jakub Kicinski <kuba@kernel.org>

1 files changed, 14 insertions, 1 deletions

diff --git a/tools/testing/selftests/net/lib/py/nsim.py b/tools/testing/selftests/net/lib/py/nsim.py
index 97457aca7e08..94aa32f59fdb 100644
--- a/tools/testing/selftests/net/lib/py/nsim.py
+++ b/tools/testing/selftests/net/lib/py/nsim.py
@@ -84,6 +84,17 @@ class NetdevSimDev:
         for port_index in range(port_count):
             self.nsims.append(self._make_port(port_index, ifnames[port_index]))
 
+        self.removed = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, ex_type, ex_value, ex_tb):
+        """
+        __exit__ gets called at the end of a "with" block.
+        """
+        self.remove()
+
     def _make_port(self, port_index, ifname):
         return NetdevSim(self, port_index, ifname, self.ns)
 
@@ -112,7 +123,9 @@ class NetdevSimDev:
             raise Exception("netdevices did not appear within timeout")
 
     def remove(self):
-        self.ctrl_write("del_device", "%u" % (self.addr, ))
+        if not self.removed:
+            self.ctrl_write("del_device", "%u" % (self.addr, ))
+            self.removed = True
 
     def remove_nsim(self, nsim):
         self.nsims.remove(nsim)