Merge branch 'Qinusty/unit-test-utils' into 'master'

Fix issue with _pretty_size with large numbers of bytes

See merge request BuildStream/buildstream!799

3 files changed, 78 insertions, 3 deletions

diff --git a/buildstream/utils.py b/buildstream/utils.py
index c116797bd..dc66f3e62 100644
--- a/buildstream/utils.py
+++ b/buildstream/utils.py
@@ -634,7 +634,7 @@ def _parse_size(size, volume):
 
 # _pretty_size()
 #
-# Converts a number of bytes into a string representation in KB, MB, GB, TB
+# Converts a number of bytes into a string representation in KiB, MiB, GiB, TiB
 # represented as K, M, G, T etc.
 #
 # Args:
@@ -646,10 +646,11 @@ def _parse_size(size, volume):
 def _pretty_size(size, dec_places=0):
     psize = size
     unit = 'B'
-    for unit in ('B', 'K', 'M', 'G', 'T'):
+    units = ('B', 'K', 'M', 'G', 'T')
+    for unit in units:
         if psize < 1024:
             break
-        else:
+        elif unit != units[-1]:
             psize /= 1024
     return "{size:g}{unit}".format(size=round(psize, dec_places), unit=unit)
 
diff --git a/tests/testutils/mock_os.py b/tests/testutils/mock_os.py
new file mode 100644
index 000000000..109593b16
--- /dev/null
+++ b/tests/testutils/mock_os.py
@@ -0,0 +1,44 @@
+from contextlib import contextmanager
+import os
+
+
+# MockAttributeResult
+#
+# A class to take a dictionary of kwargs and make them accessible via
+# attributes of the object.
+#
+class MockAttributeResult(dict):
+    __getattr__ = dict.get
+
+
+# mock_statvfs():
+#
+# Gets a function which mocks statvfs and returns a statvfs result with the kwargs accessible.
+#
+# Returns:
+#    func(path) -> object: object will have all the kwargs accessible via object.kwarg
+#
+# Example:
+#    statvfs = mock_statvfs(f_blocks=10)
+#    result = statvfs("regardless/of/path")
+#    assert result.f_blocks == 10 # True
+def mock_statvfs(**kwargs):
+    def statvfs(path):
+        return MockAttributeResult(kwargs)
+    return statvfs
+
+
+# monkey_patch()
+#
+# with monkey_patch("statvfs", custom_statvfs):
+#    assert os.statvfs == custom_statvfs # True
+# assert os.statvfs == custom_statvfs # False
+#
+@contextmanager
+def monkey_patch(to_patch, patched_func):
+    orig = getattr(os, to_patch)
+    setattr(os, to_patch, patched_func)
+    try:
+        yield
+    finally:
+        setattr(os, to_patch, orig)
diff --git a/tests/utils/misc.py b/tests/utils/misc.py
new file mode 100644
index 000000000..ae584e4d5
--- /dev/null
+++ b/tests/utils/misc.py
@@ -0,0 +1,30 @@
+from buildstream import _yaml
+from ..testutils import mock_os
+from ..testutils.runcli import cli
+
+import os
+import pytest
+
+
+KiB = 1024
+MiB = (KiB * 1024)
+GiB = (MiB * 1024)
+TiB = (GiB * 1024)
+
+
+def test_parse_size_over_1024T(cli, tmpdir):
+    BLOCK_SIZE = 4096
+    cli.configure({
+        'cache': {
+            'quota': 2048 * TiB
+        }
+    })
+    project = tmpdir.join("main")
+    os.makedirs(str(project))
+    _yaml.dump({'name': 'main'}, str(project.join("project.conf")))
+
+    bavail = (1025 * TiB) / BLOCK_SIZE
+    patched_statvfs = mock_os.mock_statvfs(f_bavail=bavail, f_bsize=BLOCK_SIZE)
+    with mock_os.monkey_patch("statvfs", patched_statvfs):
+        result = cli.run(project, args=["build", "file.bst"])
+        assert "1025T of available system storage" in result.stderr