mpm: Make contextmanagers exception-safe

When making context managers in Python, the yield statement has to be wrapped in a try/finally clause in order to properly clean up after exceptions happen.

4 files changed, 68 insertions, 5 deletions

diff --git a/mpm/python/tests/mpm_utils_tests.py b/mpm/python/tests/mpm_utils_tests.py
new file mode 100644
index 000000000..f32755a95
--- /dev/null
+++ b/mpm/python/tests/mpm_utils_tests.py
@@ -0,0 +1,55 @@
+#
+# Copyright 2020 Ettus Research, a National Instruments Brand
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+#
+import unittest
+from base_tests import TestBase
+from usrp_mpm import mpmutils
+
+
+class MockLockable:
+    """
+    Class which exposes whether lock() or unlock() have been called on it
+    """
+    def __init__(self):
+        self.locked = False
+
+    def lock(self):
+        self.locked = True
+
+    def unlock(self):
+        self.locked = False
+
+
+class TestMpmUtils(TestBase):
+    """
+    Tests for the myriad utilities in mpmutils
+    """
+    def test_normal_usage(self):
+        """
+        Checks whether in normal operation the resource gets unlocked
+        """
+        my_resource = MockLockable()
+        with mpmutils.lock_guard(my_resource):
+            self.assertEqual(my_resource.locked, True)
+        self.assertEqual(my_resource.locked, False)
+
+    def test_unlocks_after_exception(self):
+        """
+        Checked whether the resource gets unlocked after an exception occurs
+        """
+        my_resource = MockLockable()
+        try:
+            with mpmutils.lock_guard(my_resource):
+                self.assertEqual(my_resource.locked, True)
+                raise Exception("This is just a drill")
+        except Exception:
+            # Eat the raised exception
+            pass
+        finally:
+            self.assertEqual(my_resource.locked, False)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/mpm/python/tests/run_unit_tests.py b/mpm/python/tests/run_unit_tests.py
index 26fc0e1fb..88ea1a805 100755
--- a/mpm/python/tests/run_unit_tests.py
+++ b/mpm/python/tests/run_unit_tests.py
@@ -11,13 +11,17 @@ import unittest
 import sys
 import argparse
 from sys_utils_tests import TestNet
+from mpm_utils_tests import TestMpmUtils
 
 import importlib.util
 if importlib.util.find_spec("xmlrunner"):
     from xmlrunner import XMLTestRunner
 
 TESTS = {
-    '__all__': {TestNet},
+    '__all__': {
+        TestNet,
+        TestMpmUtils,
+    },
     'n3xx': set(),
 }
 
diff --git a/mpm/python/usrp_mpm/mpmutils.py b/mpm/python/usrp_mpm/mpmutils.py
index a0716d1da..f7e4e3be1 100644
--- a/mpm/python/usrp_mpm/mpmutils.py
+++ b/mpm/python/usrp_mpm/mpmutils.py
@@ -182,6 +182,8 @@ def lock_guard(lockable):
     lockable -- Must have a .lock() and .unlock() method
     """
     lockable.lock()
-    yield
-    lockable.unlock()
+    try:
+        yield
+    finally:
+        lockable.unlock()
 
diff --git a/mpm/python/usrp_mpm/sys_utils/uio.py b/mpm/python/usrp_mpm/sys_utils/uio.py
index c724557e6..84e4b2b64 100644
--- a/mpm/python/usrp_mpm/sys_utils/uio.py
+++ b/mpm/python/usrp_mpm/sys_utils/uio.py
@@ -24,8 +24,10 @@ def open_uio(label=None, path=None, length=None, read_only=True, offset=None):
     Use this like you would open() for a file"""
     uio_obj = UIO(label, path, length, read_only, offset)
     uio_obj._open()
-    yield uio_obj
-    uio_obj._close()
+    try:
+        yield uio_obj
+    finally:
+        uio_obj._close()
 
 
 def get_all_uio_devs():