Fix get_partition_by_mountpoint to actually return non '/' mount points

Change-Id: Ia11f37eef9aa927414f8b487eba3a276cc91b2aa

1 files changed, 6 insertions, 3 deletions

diff --git a/extensions/pyfdisk.py b/extensions/pyfdisk.py
index 265dc64e..a7796729 100644
--- a/extensions/pyfdisk.py
+++ b/extensions/pyfdisk.py
@@ -464,9 +464,12 @@ class Device(object):
     def get_partition_by_mountpoint(self, mountpoint):
         """Return a Partition with a specified mountpoint"""
 
-        return next(r for r in self.partitionlist
-            if hasattr(r, 'mountpoint')
-            and r.mountpoint == '/')
+        try:
+            return next(r for r in self.partitionlist
+                if hasattr(r, 'mountpoint')
+                and r.mountpoint == mountpoint)
+        except StopIteration:
+            return False
 
     def commit(self):
         """Write the partition table to the disk or image"""