Refactored shutdown/reboot tests

Tweaks to shutdownmanager to use deferLater.

landscape/client/manager/shutdownmanager.py

@@ -1,7 +1,9 @@
 import logging
-import dbus
 
+import dbus
 from twisted.internet import reactor
+from twisted.internet import task
+
 from landscape.client.manager.plugin import FAILED
 from landscape.client.manager.plugin import ManagerPlugin
 from landscape.client.manager.plugin import SUCCEEDED
@@ -21,54 +23,60 @@ class ShutdownManager(ManagerPlugin):
     This is usually sufficent.
     """
 
-    callLater = reactor.callLater
-    shutdown_delay = 120
+    def __init__(self, dbus=dbus, shutdown_delay=120):
+        self.dbus_sysbus = dbus.SystemBus()
+        self.shutdown_delay = shutdown_delay
 
     def register(self, registry):
         super().register(registry)
         self.config = registry.config
 
         registry.register_message("shutdown", self._handle_shutdown)
 
-    def _handle_shutdown(self, message):
+    def _handle_shutdown(self, message, DBus_System_Bus=None):
         """
         Choose shutdown or reboot
         """
         operation_id = message["operation-id"]
         reboot = message["reboot"]
 
-        if (reboot):
+        if reboot:
             logging.info("Reboot Requested")
             deferred = self._respond_reboot_success(
                 "Reboot requested of the system",
-                operation_id)
+                operation_id,
+            )
             return deferred
         else:
             logging.info("Shutdown Requested")
             deferred = self._respond_shutdown_success(
                 "Shutdown requested of the system",
-                operation_id)
+                operation_id,
+            )
             return deferred
 
-    def _Reboot(self, _):
+    def _Reboot(self, _, Dbus_System_bus=None):
         logging.info("Sending Reboot Command")
-        bus = dbus.SystemBus()
-        bus_object = bus.get_object(
+
+        bus_object = self.dbus_sysbus.get_object(
             "org.freedesktop.login1",
-            "/org/freedesktop/login1")
+            "/org/freedesktop/login1",
+        )
         bus_object.Reboot(
             True,
-            dbus_interface="org.freedesktop.login1.Manager")
+            dbus_interface="org.freedesktop.login1.Manager",
+        )
 
     def _Shutdown(self):
         logging.info("Sending Shutdown Command")
-        bus = dbus.SystemBus()
-        bus_object = bus.get_object(
+        bus_object = self.dbus_sysbus.get_object(
             "org.freedesktop.login1",
-            "/org/freedesktop/login1")
+            "/org/freedesktop/login1",
+        )
         bus_object.PowerOff(
             True,
-            dbus_interface="org.freedesktop.login1.Manager")
+            dbus_interface="org.freedesktop.login1.Manager",
+        )
 
     def _respond_reboot_success(self, data, operation_id):
         deferred = self._respond(SUCCEEDED, data, operation_id)
@@ -78,7 +86,11 @@ def _respond_reboot_success(self, data, operation_id):
 
     def _respond_shutdown_success(self, data, operation_id):
         deferred = self._respond(SUCCEEDED, data, operation_id)
-        self.callLater(self.shutdown_delay, self._Shutdown)
+        self.shutdown_deferred = task.deferLater(
+            reactor,
+            self.shutdown_delay,
+            self._Shutdown,
+        )
         deferred.addErrback(self._respond_fail)
         return deferred
 

landscape/client/manager/tests/test_shutdownmanager.py

@@ -1,5 +1,4 @@
-from unittest.mock import patch, Mock
-from twisted.internet import task
+from unittest.mock import Mock
 
 from landscape.client.manager.shutdownmanager import ShutdownManager
 from landscape.client.tests.helpers import LandscapeTest
@@ -18,39 +17,30 @@ def setUp(self):
         )
         self.broker_service.pinger.start()
 
-        self.clock = task.Clock()
-        self.plugin = ShutdownManager()
+        self.dbus_mock = Mock()
+        self.dbus_sysbus_mock = Mock()
+        self.dbus_mock.get_object.return_value = self.dbus_sysbus_mock
+        self.plugin = ShutdownManager(dbus=self.dbus_mock, shutdown_delay=0)
         self.manager.add(self.plugin)
 
-        self.dbus_mock = patch(
-            "landscape.client.manager.shutdownmanager.dbus").start()
-
     def test_reboot(self):
-        bus_object = Mock()
-        self.dbus_mock.SystemBus.return_value = bus_object
-
         message = {"type": "shutdown", "reboot": True, "operation-id": 100}
         deferred = self.plugin._handle_shutdown(message)
 
         def check(_):
-            bus_object.get_object.assert_called_once()
-            bus_object.Reboot.assert_called_once()
+            self.plugin.dbus_sysbus.get_object.assert_called_once()
+            self.plugin.dbus_sysbus.get_object().Reboot.assert_called_once()
 
         deferred.addCallback(check)
         return deferred
 
     def test_shutdown(self):
-        bus_object = Mock()
-        self.dbus_mock.SystemBus.return_value = bus_object
-
         message = {"type": "shutdown", "reboot": False, "operation-id": 100}
         deferred = self.plugin._handle_shutdown(message)
 
-        self.clock.advance(self.plugin.shutdown_delay)
-
         def check(_):
-            bus_object.get_object.assert_called_once()
-            bus_object.PowerOff.assert_called_once()
+            self.plugin.dbus_sysbus.get_object.assert_called_once()
+            self.plugin.dbus_sysbus.get_object().PowerOff.assert_called_once()
 
-        deferred.addCallback(check)
+        self.plugin.shutdown_deferred.addCallback(check)
         return deferred