From 8554d6e39b6ad967e1debe98550a0c56aaf8c8ea Mon Sep 17 00:00:00 2001
From: Jakub Kicinski <kuba@kernel.org>
Date: Fri, 12 Apr 2024 07:14:35 -0700
Subject: [PATCH] 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>
---
 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 @@ 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)