Added directory_usage to utils (#629)

* Added console_has_banner parameter This parameter allows to modify login state machine behavior in case of a hung state. * Added directory_usage to utils This utility is similar to "show system directory-usage" Junos command (or UNIX "du" command). It allows to see how much space is taken by a particular directory on disk. Unit test added as well. * Modified directory_usage() to return dict and add depth parameter * Improving exception handling in directory_usage() util * Update test_fs.py
Juniper · Dec 5, 2016 · 6b9faac · 6b9faac
1 parent b84b8a6
commit 6b9faac
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 2 deletions.
diff --git a/lib/jnpr/junos/utils/fs.py b/lib/jnpr/junos/utils/fs.py
@@ -3,6 +3,8 @@
 from jnpr.junos.utils.util import Util
 from jnpr.junos.utils.start_shell import StartShell
 
+from jnpr.junos.exception import RpcError
+
 
 class FS(Util):
     """
@@ -20,6 +22,7 @@ class FS(Util):
     * :meth:`rmdir`: remove a directory
     * :meth:`stat`: return file/dir information
     * :meth:`storage_usage`: return storage usage
+    * :meth:`directory_usage`: return directory usage
     * :meth:`storage_cleanup`: perform storage storage_cleanup
     * :meth:`storage_cleanup_check`: returns a list of files to remove at cleanup
     * :meth:`symlink`: create a symlink
@@ -263,6 +266,42 @@ def _decode(fs):
 
         return dict((_name(fs), _decode(fs)) for fs in rsp.xpath('filesystem'))
 
+    # -------------------------------------------------------------------------
+    # directory_usage - filesystem directory usage
+    # -------------------------------------------------------------------------
+
+    def directory_usage(self, path=".", depth=0):
+        """
+        Returns the directory usage, similar to the unix "du" command.
+
+        :returns: dict of directory usage, including subdirectories if depth > 0
+        """
+        BLOCK_SIZE = 512
+
+        rsp = self._dev.rpc.get_directory_usage_information(path=path, depth=str(depth))
+
+        result = {}
+
+        for directory in rsp.findall(".//directory"):
+            dir_name = directory.findtext("directory-name").strip()
+            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
+
     # -------------------------------------------------------------------------
     ### storage_cleanup_check, storage_cleanip
     # -------------------------------------------------------------------------
@@ -388,8 +427,8 @@ def tgz(self, from_path, tgz_path):
         return rsp.text
 
     # -------------------------------------------------------------------------
-    # !!!!! methods that use SSH shell commands, requires that the user
-    # !!!!! has 'start shell' priveldges
+    # !!!!! methods that use SSH shell commands, require that the user
+    # !!!!! has 'start shell' privileges
     # -------------------------------------------------------------------------
 
     def _ssh_exec(self, command):

diff --git a/tests/unit/utils/rpc-reply/get-directory-usage-information.xml b/tests/unit/utils/rpc-reply/get-directory-usage-information.xml
@@ -0,0 +1,49 @@
+<rpc-reply xmlns:junos="http://xml.juniper.net/junos/12.1X47/junos">
+    <directory-usage-information>
+        <directory>
+            <directory-name>/var/tmp</directory-name>
+            <directory>
+                <directory-name>/var/tmp/install</directory-name>
+                <used-space used-blocks="4">
+                    2.0K
+                </used-space>
+            </directory>
+            <directory>
+                <directory-name>/var/tmp/vi.recover</directory-name>
+                <used-space used-blocks="4">
+                    2.0K
+                </used-space>
+            </directory>
+            <directory>
+                <directory-name>/var/tmp/pics</directory-name>
+                <used-space used-blocks="4">
+                    2.0K
+                </used-space>
+            </directory>
+            <directory>
+                <directory-name>/var/tmp/gres-tp</directory-name>
+                <used-space used-blocks="68">
+                    34K
+                </used-space>
+            </directory>
+            <directory>
+                <directory-name>/var/tmp/rtsdb</directory-name>
+                <used-space used-blocks="4">
+                    2.0K
+                </used-space>
+            </directory>
+            <directory>
+                <directory-name>/var/tmp/sec-download</directory-name>
+                <used-space used-blocks="8">
+                    4.0K
+                </used-space>
+            </directory>
+            <used-space used-blocks="456076">
+                223M
+            </used-space>
+        </directory>
+    </directory-usage-information>
+    <cli>
+        <banner></banner>
+    </cli>
+</rpc-reply>
diff --git a/tests/unit/utils/test_fs.py b/tests/unit/utils/test_fs.py
@@ -266,6 +266,19 @@ def test_storage_usage(self, mock_execute):
                            'avail': '2F', 'used': '481M',
                            'total': '4F'}})
 
+    @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", 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'},
+                          '/var/tmp/pics': {'blocks': 4, 'bytes': 2048, 'size': '2.0K'},
+                          '/var/tmp/rtsdb': {'blocks': 4, 'bytes': 2048, 'size': '2.0K'},
+                          '/var/tmp/sec-download': {'blocks': 8, 'bytes': 4096, 'size': '4.0K'},
+                          '/var/tmp/vi.recover': {'blocks': 4, 'bytes': 2048, 'size': '2.0K'}}
+                         )
+
     @patch('jnpr.junos.Device.execute')
     def test_storage_cleanup(self, mock_execute):
         mock_execute.side_effect = self._mock_manager
@@ -332,6 +345,8 @@ def _mock_manager(self, *args, **kwargs):
                     return self._read_file('show-cli-directory.xml')
             elif args[0].tag == 'get-system-storage':
                 return self._read_file('get-system-storage.xml')
+            elif args[0].tag == 'get-directory-usage-information':
+                return self._read_file('get-directory-usage-information.xml')
             elif args[0].tag == 'request-system-storage-cleanup':
                 return self._read_file('request-system-storage-cleanup.xml')
             elif args[0].tag == 'file-archive':