Skip to content
Open
22 changes: 11 additions & 11 deletions salt/modules/disk.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,19 +39,19 @@ def _parse_numbers(text):
return decimal.Decimal(text)

try:
postPrefixes = {
"K": "10E3",
"M": "10E6",
"G": "10E9",
"T": "10E12",
"P": "10E15",
"E": "10E18",
"Z": "10E21",
"Y": "10E24",
unit_suffixes = {
"K": "1E3",
"M": "1E6",
"G": "1E9",
"T": "1E12",
"P": "1E15",
"E": "1E18",
"Z": "1E21",
"Y": "1E24",
}
if text[-1] in postPrefixes.keys():
if text[-1] in unit_suffixes.keys():
v = decimal.Decimal(text[:-1])
v = v * decimal.Decimal(postPrefixes[text[-1]])
v = v * decimal.Decimal(unit_suffixes[text[-1]])
return v
else:
return decimal.Decimal(text)
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/modules/test_disk.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
import unittest
from salt.modules.disk import _parse_numbers

class TestParseNumbers(unittest.TestCase):

def test_integer(self):
result = _parse_numbers("42")
self.assertEqual(result, decimal.Decimal("42"))

def test_with_kilo(self):
result = _parse_numbers("1.5K")
self.assertEqual(result, decimal.Decimal("1500"))

def test_with_mega(self):
result = _parse_numbers("2.5M")
self.assertEqual(result, decimal.Decimal("2500000"))

# Add more test cases for other units and scenarios

if __name__ == '__main__':
unittest.main()