diff --git a/bumblebee/modules/traffic.py b/bumblebee/modules/traffic.py index 250a1a7..856ac76 100644 --- a/bumblebee/modules/traffic.py +++ b/bumblebee/modules/traffic.py @@ -10,6 +10,7 @@ Parameters: Defaults to "{:.2f}" """ +import re import time import psutil import netifaces @@ -70,6 +71,20 @@ class Module(bumblebee.engine.Module): return [] return retval + def get_minwidth_str(self): + """computes theme.minwidth string based on traffic.format parameter""" + minwidth_str = "1000" + try: + length = int(re.match("{:\.(\d+)f}", self._format).group(1)) + if length > 0: + minwidth_str += "." + "0" * length + except AttributeError: + # return default value + return "1000.00MB" + finally: + minwidth_str += "MB" + return minwidth_str + def _update_widgets(self, widgets): interfaces = [i for i in netifaces.interfaces() if not i.startswith(self._exclude)] @@ -103,7 +118,7 @@ class Module(bumblebee.engine.Module): for direction in ["rx", "tx"]: name = "traffic.{}-{}".format(direction, interface) - widget = self.create_widget(widgets, name, attributes={"theme.minwidth": "1000.00MB"}) + widget = self.create_widget(widgets, name, attributes={"theme.minwidth": self.get_minwidth_str()}) prev = self._prev.get(name, 0) speed = bumblebee.util.bytefmt( (int(data[direction]) - int(prev))/timediff, diff --git a/tests/modules/test_traffic.py b/tests/modules/test_traffic.py new file mode 100644 index 0000000..951a4d8 --- /dev/null +++ b/tests/modules/test_traffic.py @@ -0,0 +1,23 @@ +import mock +import unittest + +import tests.mocks as mocks + +from bumblebee.modules.traffic import Module + +class TestTrafficModule(unittest.TestCase): + def setUp(self): + mocks.setup_test(self, Module) + + def test_default_format(self): + self.assertEqual(self.module._format, "{:.2f}") + + def test_get_minwidth_str(self): + # default value (two digits after dot) + self.assertEqual(self.module.get_minwidth_str(), "1000.00MB") + # integer value + self.module._format = "{:.0f}" + self.assertEqual(self.module.get_minwidth_str(), "1000MB") + # just one digit after dot + self.module._format = "{:.1f}" + self.assertEqual(self.module.get_minwidth_str(), "1000.0MB")