Fix slurm_format_bytes_ceil logic and add unit tests (#675)

* Fix slurm_format_bytes_ceil logic, add unit tests

* Set default return value to 1K

Author: maneesh29s

dask_jobqueue/slurm.py

@@ -116,13 +116,26 @@ def slurm_format_bytes_ceil(n):
     """Format bytes as text.
 
     SLURM expects KiB, MiB or Gib, but names it KB, MB, GB. SLURM does not handle Bytes, only starts at KB.
+    Thus minimum returned value is '1K'.
 
-    >>> slurm_format_bytes_ceil(1)
+    Parameters
+    ----------
+    n: int
+        Number of bytes
+
+    Returns
+    -------
+    str
+        Bytes formatted as string
+
+    Examples
+    --------
+    >>> slurm_format_bytes_ceil(12)
     '1K'
     >>> slurm_format_bytes_ceil(1234)
     '2K'
     >>> slurm_format_bytes_ceil(12345678)
-    '13M'
+    '12M'
     >>> slurm_format_bytes_ceil(1234567890)
     '2G'
     >>> slurm_format_bytes_ceil(15000000000)
@@ -134,7 +147,7 @@ def slurm_format_bytes_ceil(n):
         return "%dM" % math.ceil(n / (1024**2))
     if n >= 1024:
         return "%dK" % math.ceil(n / 1024)
-    return "1K" % n
+    return "1K"
 
 
 class SLURMCluster(JobQueueCluster):

dask_jobqueue/tests/test_slurm.py

@@ -10,6 +10,7 @@
 from dask.utils import format_bytes, parse_bytes
 
 from dask_jobqueue import SLURMCluster
+from dask_jobqueue.slurm import slurm_format_bytes_ceil
 
 from . import QUEUE_WAIT
 
@@ -343,3 +344,18 @@ def test_slurm_runner():
     )
     output = output.decode()
     assert "Test passed" in output
+
+
+@pytest.mark.parametrize(
+    ["input", "expected_output"],
+    [
+        (12, "1K"),
+        (1234, "2K"),
+        (12345678, "12M"),
+        (15000000000, "14G"),
+    ],
+)
+def test_slurm_format_bytes_ceil(input, expected_output):
+    actual_output = slurm_format_bytes_ceil(input)
+
+    assert actual_output == expected_output