From 225d471c6a4a7cc93b5049c9e633665e74054f50 Mon Sep 17 00:00:00 2001
From: Tobi-wan Kenobi <github@tobi-wan-kenobi.at>
Date: Sat, 10 Dec 2016 08:09:13 +0100
Subject: [PATCH] [modules/cpu] Add configurable warning and critical
 thresholds

The cpu module now has cpu.warning and cpu.critical thresholds. If the
CPU utilization is higher than any of those values, the widget's state
changes to warning or critical, respectively.

see #23
---
 bumblebee-status          | 26 +++++++++++++-------------
 bumblebee/modules/cpu.py  | 14 +++++++++++++-
 bumblebee/output.py       |  6 +++++-
 tests/modules/test_cpu.py | 23 +++++++++++++++++++++--
 tests/util.py             |  8 ++++++++
 5 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/bumblebee-status b/bumblebee-status
index 44eae21..248ff52 100755
--- a/bumblebee-status
+++ b/bumblebee-status
@@ -18,19 +18,19 @@ def main():
         inp=inp,
     )
 
-    try:
-        engine.run()
-    except KeyboardInterrupt as error:
-        inp.stop()
-        sys.exit(0)
-    except bumblebee.error.BaseError as error:
-        inp.stop()
-        sys.stderr.write("fatal: {}\n".format(error))
-        sys.exit(1)
-    except Exception as error:
-        inp.stop()
-        sys.stderr.write("fatal: {}\n".format(error))
-        sys.exit(2)
+    engine.run()
+#    try:
+#    except KeyboardInterrupt as error:
+#        inp.stop()
+#        sys.exit(0)
+#    except bumblebee.error.BaseError as error:
+#        inp.stop()
+#        sys.stderr.write("fatal: {}\n".format(error))
+#        sys.exit(1)
+#    except Exception as error:
+#        inp.stop()
+#        sys.stderr.write("fatal: {}\n".format(error))
+#        sys.exit(2)
 
 if __name__ == "__main__":
     main()
diff --git a/bumblebee/modules/cpu.py b/bumblebee/modules/cpu.py
index 3d20479..bf4cbc6 100644
--- a/bumblebee/modules/cpu.py
+++ b/bumblebee/modules/cpu.py
@@ -1,6 +1,11 @@
 # pylint: disable=C0111,R0903
 
-"""Displays CPU utilization across all CPUs."""
+"""Displays CPU utilization across all CPUs.
+
+Parameters:
+    * cpu.warning : Warning threshold in % of CPU usage (defaults to 70%)
+    * cpu.critical: Critical threshold in % of CPU usage (defaults to 80%)
+"""
 
 import psutil
 import bumblebee.input
@@ -22,4 +27,11 @@ class Module(bumblebee.engine.Module):
     def update(self, widgets):
         self._utilization = psutil.cpu_percent(percpu=False)
 
+    def state(self, widget):
+        if self._utilization > int(self.parameter("critical", 80)):
+            return "critical"
+        if self._utilization > int(self.parameter("warning", 70)):
+            return "warning"
+        return None
+
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
diff --git a/bumblebee/output.py b/bumblebee/output.py
index 30062f3..70bd864 100644
--- a/bumblebee/output.py
+++ b/bumblebee/output.py
@@ -11,6 +11,7 @@ class Widget(object):
     def __init__(self, full_text="", name=""):
         self._full_text = full_text
         self.module = None
+        self._module = None
         self.name = name
         self.id = str(uuid.uuid4())
 
@@ -20,10 +21,13 @@ class Widget(object):
         This is done outside the constructor to avoid having to
         pass in the module name in every concrete module implementation"""
         self.module = module.name
+        self._module = module
 
     def state(self):
         """Return the widget's state"""
-        return "state-default"
+        if self._module and hasattr(self._module, "state"):
+            return self._module.state(self)
+        return None
 
     def full_text(self):
         """Retrieve the full text to display in the widget"""
diff --git a/tests/modules/test_cpu.py b/tests/modules/test_cpu.py
index 4bce1e2..241656c 100644
--- a/tests/modules/test_cpu.py
+++ b/tests/modules/test_cpu.py
@@ -7,14 +7,17 @@ import mock
 import bumblebee.input
 from bumblebee.input import I3BarInput
 from bumblebee.modules.cpu import Module
-from tests.util import MockEngine, assertPopen
+from tests.util import MockEngine, MockConfig, assertPopen
 
 class TestCPUModule(unittest.TestCase):
     def setUp(self):
         self.engine = MockEngine()
         self.engine.input = I3BarInput()
         self.engine.input.need_event = True
-        self.module = Module(engine=self.engine, config={})
+        self.config = MockConfig()
+        self.module = Module(engine=self.engine, config={ "config": self.config })
+        for widget in self.module.widgets():
+            widget.link_module(self.module)
 
     @mock.patch("sys.stdout")
     def test_format(self, mock_output):
@@ -34,4 +37,20 @@ class TestCPUModule(unittest.TestCase):
         mock_input.readline.assert_any_call()
         assertPopen(mock_output, "gnome-system-monitor")
 
+    @mock.patch("psutil.cpu_percent")
+    def test_warning(self, mock_psutil):
+        self.config.set("cpu.critical", "20")
+        self.config.set("cpu.warning", "18")
+        mock_psutil.return_value = 19.0
+        self.module.update(self.module.widgets())
+        self.assertEquals(self.module.widgets()[0].state(), "warning")
+
+    @mock.patch("psutil.cpu_percent")
+    def test_critical(self, mock_psutil):
+        self.config.set("cpu.critical", "20")
+        self.config.set("cpu.warning", "19")
+        mock_psutil.return_value = 21.0
+        self.module.update(self.module.widgets())
+        self.assertEquals(self.module.widgets()[0].state(), "critical")
+
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
diff --git a/tests/util.py b/tests/util.py
index 6b04725..f755a7a 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -31,9 +31,17 @@ class MockEngine(object):
         self.input = MockInput()
 
 class MockConfig(object):
+    def __init__(self):
+        self._data = {}
+
     def get(self, name, default):
+        if name in self._data:
+            return self._data[name]
         return default
 
+    def set(self, name, value):
+        self._data[name] = value
+
 class MockOutput(object):
     def start(self):
         pass