[PATCH net-next v2 5/6] selftests: net: support use of NetdevSimDev under "with" in python

[Jakub Kicinski <kuba@xxxxxxxxxx>]

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@xxxxxxxxxx>
Signed-off-by: Jakub Kicinski <kuba@xxxxxxxxxx>
---
 tools/testing/selftests/net/lib/py/nsim.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

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 @@ from .utils import cmd, ip
         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 @@ from .utils import cmd, ip
             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)
-- 
2.44.0