Support "with"-style errno assertions. am: 2f4685bfae am: 8159f234a1 am: b3a5dd6764oreo-dr1-dev

am: 3fa70dadad

Change-Id: I1f1c6388109b6ab69c29a663b134e43e3e1a5626

1 files changed, 20 insertions, 2 deletions

diff --git a/net/test/net_test.py b/net/test/net_test.py
index 9a5c635e..423228be 100755
--- a/net/test/net_test.py
+++ b/net/test/net_test.py
@@ -375,9 +375,27 @@ class RunAsUid(RunAsUidGid):
 
 class NetworkTest(unittest.TestCase):
 
-  def assertRaisesErrno(self, err_num, f, *args):
+  def assertRaisesErrno(self, err_num, f=None, *args):
+    """Test that the system returns an errno error.
+
+    This works similarly to unittest.TestCase.assertRaises. You can call it as
+    an assertion, or use it as a context manager.
+    e.g.
+        self.assertRaisesErrno(errno.ENOENT, do_things, arg1, arg2)
+    or
+        with self.assertRaisesErrno(errno.ENOENT):
+          do_things(arg1, arg2)
+
+    Args:
+      err_num: an errno constant
+      f: (optional) A callable that should result in error
+      *args: arguments passed to f
+    """
     msg = os.strerror(err_num)
-    self.assertRaisesRegexp(EnvironmentError, msg, f, *args)
+    if f is None:
+      return self.assertRaisesRegexp(EnvironmentError, msg)
+    else:
+      self.assertRaisesRegexp(EnvironmentError, msg, f, *args)
 
   def ReadProcNetSocket(self, protocol):
     # Read file.