From e15147fe107e53f113c3d4ab84d7aab0717c3376 Mon Sep 17 00:00:00 2001
From: Tobi-wan Kenobi <github@tobi-wan-kenobi.at>
Date: Sat, 10 Dec 2016 18:49:25 +0100
Subject: [PATCH] [tests/disk] Add critical/warning threshold tests for disk
 module

see #23
---
 bumblebee/modules/disk.py  | 25 +++++++++++--------------
 tests/modules/test_disk.py | 25 +++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/bumblebee/modules/disk.py b/bumblebee/modules/disk.py
index 586909f..a0b6aad 100644
--- a/bumblebee/modules/disk.py
+++ b/bumblebee/modules/disk.py
@@ -26,25 +26,22 @@ class Module(bumblebee.engine.Module):
             cmd="nautilus {}".format(self._path))
 
     def diskspace(self):
-        st = os.statvfs(self._path)
-        size = st.f_frsize*st.f_blocks
-        used = size - st.f_frsize*st.f_bavail
-        self._perc = 100.0*used/size
-
         return "{} {}/{} ({:05.02f}%)".format(self._path,
-            bumblebee.util.bytefmt(used),
-            bumblebee.util.bytefmt(size), self._perc
+            bumblebee.util.bytefmt(self._used),
+            bumblebee.util.bytefmt(self._size), self._perc
         )
 
     def update(self, widgets):
-        pass
+        st = os.statvfs(self._path)
+        self._size = st.f_frsize*st.f_blocks
+        self._used = self._size - st.f_frsize*st.f_bavail
+        self._perc = 100.0*self._used/self._size
 
     def state(self, widget):
-        pass
-    def warning(self, widget):
-        return self._perc > self._config.parameter("warning", 80)
-
-    def critical(self, widget):
-        return self._perc > self._config.parameter("critical", 90)
+        if self._perc > float(self.parameter("critical", 90)):
+            return "critical"
+        if self._perc > float(self.parameter("warning", 80)):
+            return "warning"
+        return None
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
diff --git a/tests/modules/test_disk.py b/tests/modules/test_disk.py
index 3ee74b0..4d3909c 100644
--- a/tests/modules/test_disk.py
+++ b/tests/modules/test_disk.py
@@ -9,6 +9,12 @@ from bumblebee.input import I3BarInput
 from bumblebee.modules.disk import Module
 from tests.util import MockEngine, MockConfig, assertPopen
 
+class MockVFS(object):
+    def __init__(self, perc):
+        self.f_blocks = 1024*1024
+        self.f_frsize = 1
+        self.f_bavail = self.f_blocks - self.f_blocks*(perc/100.0)
+
 class TestDiskModule(unittest.TestCase):
     def setUp(self):
         self.engine = MockEngine()
@@ -17,6 +23,8 @@ class TestDiskModule(unittest.TestCase):
         self.config = MockConfig()
         self.config.set("disk.path", "somepath")
         self.module = Module(engine=self.engine, config={"config": self.config})
+        for widget in self.module.widgets():
+            widget.link_module(self.module)
 
     @mock.patch("select.select")
     @mock.patch("subprocess.Popen")
@@ -33,4 +41,21 @@ class TestDiskModule(unittest.TestCase):
         mock_input.readline.assert_any_call()
         assertPopen(mock_output, "nautilus {}".format(self.module.parameter("path")))
 
+    @mock.patch("os.statvfs")
+    def test_warning(self, mock_stat):
+        self.config.set("disk.critical", "80")
+        self.config.set("disk.warning", "70")
+        mock_stat.return_value = MockVFS(75.0)
+        self.module.update(self.module.widgets())
+        self.assertEquals(self.module.widgets()[0].state(), ["warning"])
+
+    @mock.patch("os.statvfs")
+    def test_warning(self, mock_stat):
+        self.config.set("disk.critical", "80")
+        self.config.set("disk.warning", "70")
+        mock_stat.return_value = MockVFS(85.0)
+        self.module.update(self.module.widgets())
+        self.assertEquals(self.module.widgets()[0].state(), ["critical"])
+
+
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4