Merge pull request #435 from cddmp/master

Improve screensaver and dpms handling in caffeine module

bumblebee/modules/caffeine.py

@@ -1,12 +1,17 @@
-# pylint: disable=C0111,R0903
+#pylint: disable=C0111,R0903,W0212
 
 """Enable/disable automatic screen locking.
 
 Requires the following executables:
-    * xset
+    * xdg-screensaver
+    * xdotool
+    * xprop (as dependency for xdotool)
     * notify-send
 """
 
+import logging
+import os
+import psutil
 import bumblebee.input
 import bumblebee.output
 import bumblebee.engine
@@ -14,35 +19,84 @@ import bumblebee.engine
 class Module(bumblebee.engine.Module):
     def __init__(self, engine, config):
         super(Module, self).__init__(engine, config,
-            bumblebee.output.Widget(full_text=self.caffeine)
+            bumblebee.output.Widget(full_text="")
         )
+        self._active = False
+        self._xid = None
+
         engine.input.register_callback(self, button=bumblebee.input.LEFT_MOUSE,
             cmd=self._toggle
         )
 
-    def caffeine(self, widget):
+    def _check_requirements(self):
-        return ""
+        requirements = ['xdotool', 'xprop', 'xdg-screensaver']
+        missing = []
+        for tool in requirements:
+            if not bumblebee.util.which(tool):
+                missing.append(tool)
+        return missing
 
-    def state(self, widget):
+    def _get_i3bar_xid(self):
-        if self._active():
+        xid = bumblebee.util.execute("xdotool search --class \"i3bar\"").partition('\n')[0].strip()
+        if xid.isdigit():
+            return xid
+        logging.warning("Module caffeine: xdotool couldn't get X window ID of \"i3bar\".")
+        return None
+
+    def _notify(self):
+        if not bumblebee.util.which('notify-send'):
+            return
+
+        if self._active:
+            bumblebee.util.execute("notify-send \"Consuming caffeine\"")
+        else:
+            bumblebee.util.execute("notify-send \"Out of coffee\"")
+
+    def _suspend_screensaver(self):
+        self._xid = self._get_i3bar_xid()
+        if self._xid is None:
+            return False
+
+        pid = os.fork()
+        if pid == 0:
+            os.setsid()
+            bumblebee.util.execute("xdg-screensaver suspend {}".format(self._xid))
+            os._exit(0)
+        else:
+            os.waitpid(pid, 0)
+        return True
+
+    def _resume_screensaver(self):
+        success = True
+        xprop_path = bumblebee.util.which('xprop')
+        pids = [ p.pid for p in psutil.process_iter() if p.cmdline() == [xprop_path, '-id', str(self._xid), '-spy'] ]
+        for pid in pids:
+            try:
+                os.kill(pid, 9)
+            except OSError:
+                success = False
+        return success
+
+    def state(self, _):
+        if self._active:
             return "activated"
         return "deactivated"
 
-    def _active(self):
+    def _toggle(self, _):
-        for line in bumblebee.util.execute("xset q").split("\n"):
+        missing = self._check_requirements()
-            if "timeout" in line:
+        if missing:
-                timeout = int(line.split(" ")[4])
+            logging.warning('Could not run caffeine - missing %s!', ", ".join(missing))
-                if timeout == 0:
+            return
-                    return True
-                return False
-        return False
 
-    def _toggle(self, widget):
+        self._active = not self._active
-        if self._active():
+        if self._active:
-            bumblebee.util.execute("xset s default")
+            success = self._suspend_screensaver()
-            bumblebee.util.execute("notify-send \"Out of coffee\"")
         else:
-            bumblebee.util.execute("xset s off")
+            success = self._resume_screensaver()
-            bumblebee.util.execute("notify-send \"Consuming caffeine\"")
+
+        if success:
+            self._notify()
+        else:
+            self._active = not self._active
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4

tests/modules/test_caffeine.py

@@ -1,17 +1,9 @@
 # pylint: disable=C0103,C0111
 
-import json
 import unittest
-import mock
+from mock import patch
-
-try:
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
-
 import tests.mocks as mocks
 
-from bumblebee.config import Config
 from bumblebee.input import LEFT_MOUSE
 from bumblebee.modules.caffeine import Module
 
@@ -19,35 +11,40 @@ class TestCaffeineModule(unittest.TestCase):
     def setUp(self):
         mocks.setup_test(self, Module)
 
-        self.xset_active = "  timeout:  0    cycle:  123"
-        self.xset_inactive = "  timeout:  600    cycle:  123"
-
     def tearDown(self):
         mocks.teardown_test(self)
 
-    def test_text(self):
+    def test_check_requirements(self):
-        self.assertEquals(self.module.caffeine(self.anyWidget), "")
+        with patch('bumblebee.util.which', side_effect=['', 'xprop', 'xdg-screensaver']):
+            self.assertTrue(['xdotool'] == self.module._check_requirements())
 
-    def test_active(self):
+    def test_get_i3bar_xid_returns_digit(self):
-        self.popen.mock.communicate.return_value = (self.xset_active, None)
+        self.popen.mock.communicate.return_value = ("8388614", None)
-        self.assertTrue(not "deactivated" in self.module.state(self.anyWidget))
+        self.assertTrue(self.module._get_i3bar_xid().isdigit())
-        self.assertTrue("activated" in self.module.state(self.anyWidget))
 
-    def test_inactive(self):
+    def test_get_i3bar_xid_returns_error_string(self):
-        self.popen.mock.communicate.return_value = (self.xset_inactive, None)
+        self.popen.mock.communicate.return_value = ("Some error message", None)
-        self.assertTrue("deactivated" in self.module.state(self.anyWidget))
+        self.assertTrue(self.module._get_i3bar_xid() is None)
-        self.popen.mock.communicate.return_value = ("no text", None)
-        self.assertTrue("deactivated" in self.module.state(self.anyWidget))
 
-    def test_toggle(self):
+    def test_get_i3bar_xid_returns_empty_string(self):
-        self.popen.mock.communicate.return_value = (self.xset_active, None)
+        self.popen.mock.communicate.return_value = ("", None)
+        self.assertTrue(self.module._get_i3bar_xid() is None)
+
+    def test_suspend_screensaver_success(self):
+        with patch.object(self.module, '_get_i3bar_xid', return_value=8388614):
             mocks.mouseEvent(stdin=self.stdin, button=LEFT_MOUSE, inp=self.input, module=self.module)
-        self.popen.assert_call("xset s default")
+            self.assertTrue(self.module._suspend_screensaver() is True)
-        self.popen.assert_call("notify-send \"Out of coffee\"")
 
-        self.popen.mock.communicate.return_value = (self.xset_inactive, None)
+    def test_suspend_screensaver_fail(self):
+        with patch.object(self.module, '_get_i3bar_xid', return_value=None):
+            self.module._active = False
             mocks.mouseEvent(stdin=self.stdin, button=LEFT_MOUSE, inp=self.input, module=self.module)
-        self.popen.assert_call("xset s off")
+            self.assertTrue(self.module._suspend_screensaver() is False)
-        self.popen.assert_call("notify-send \"Consuming caffeine\"")
+
+    def test_resume_screensaver(self):
+        with patch.object(self.module, '_check_requirements', return_value=[]):
+            self.module._active = True
+            mocks.mouseEvent(stdin=self.stdin, button=LEFT_MOUSE, inp=self.input, module=self.module)
+            self.assertTrue(self.module._resume_screensaver() is True)
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4