Improving exception handling in directory_usage() util

Juniper · Nov 30, 2016 · 5c80bed · 5c80bed
1 parent d2e6459
commit 5c80bed
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git a/lib/jnpr/junos/utils/fs.py b/lib/jnpr/junos/utils/fs.py
@@ -280,22 +280,25 @@ def directory_usage(self, path=".", depth=0):
 
         rsp = self._dev.rpc.get_directory_usage_information(path=path, depth=str(depth))
 
-        if rsp.findtext("directory/used-space") is None:
-            # Likely, no such directory
-            raise RpcError(rsp=rsp)
-
-        dirs = rsp.xpath("//directory")
         result = {}
 
-        for directory in dirs:
+        for directory in rsp.findall(".//directory"):
             dir_name = directory.findtext("directory-name").strip()
-            dir_size = directory.findtext("used-space").strip()
-            dir_blocks = int(directory.find("used-space").get("used-blocks").strip())
-            result[dir_name] = {
-                "size": dir_size,
-                "blocks": dir_blocks,
-                "bytes": dir_blocks * BLOCK_SIZE,
-            }
+            if dir_name is None:
+                raise RpcError(rsp=rsp)
+
+            used_space = directory.find('used-space')
+            if used_space is not None:
+                dir_size = used_space.text.strip()
+                dir_blocks = used_space.get('used-blocks')
+                if dir_blocks is not None:
+                    dir_blocks = int(dir_blocks)
+                    dir_bytes = dir_blocks * BLOCK_SIZE
+                    result[dir_name] = {
+                        "size": dir_size,
+                        "blocks": dir_blocks,
+                        "bytes": dir_bytes,
+                    }
 
         return result
 

diff --git a/tests/unit/utils/test_fs.py b/tests/unit/utils/test_fs.py
@@ -269,7 +269,7 @@ def test_storage_usage(self, mock_execute):
     @patch('jnpr.junos.Device.execute')
     def test_directory_usage(self, mock_execute):
         mock_execute.side_effect = self._mock_manager
-        self.assertEqual(self.fs.directory_usage(path="/var/tmp"),
+        self.assertEqual(self.fs.directory_usage(path="/var/tmp", depth=1),
                          {'/var/tmp': {'blocks': 456076, 'bytes': 233510912, 'size': '223M'},
                           '/var/tmp/gres-tp': {'blocks': 68, 'bytes': 34816, 'size': '34K'},
                           '/var/tmp/install': {'blocks': 4, 'bytes': 2048, 'size': '2.0K'},