diff --git a/charms_openstack/charm/core.py b/charms_openstack/charm/core.py
index 39852df..3412ee6 100644
--- a/charms_openstack/charm/core.py
+++ b/charms_openstack/charm/core.py
@@ -1,5 +1,4 @@
 import collections
-import contextlib
 import functools
 import itertools
 import os
@@ -817,28 +816,26 @@ class BaseOpenStackCharmActions(object):
         """
         ch_host.service_reload(service_name, restart_on_failure)
 
-    @contextlib.contextmanager
     def restart_on_change(self):
         """Restart the services in the self.restart_map{} attribute if any of
         the files identified by the keys changes for the wrapped call.
 
-        This function is a @decorator that checks if the wrapped function
-        changes any of the files identified by the keys in the
-        self.restart_map{} and, if they change, restarts the services in the
-        corresponding list.
+        Usage:
+
+           @restart_on_change(restart_map, ...)
+           def function_that_might_trigger_a_restart(...)
+               ...
+
+        Or:
+
+           with restart_on_change(restart_map, ...):
+               do_stuff_that_might_trigger_a_restart()
+               ...
         """
-        checksums = {path: ch_host.path_hash(path)
-                     for path in self.full_restart_map.keys()}
-        yield
-        restarts = []
-        for path in self.full_restart_map:
-            if ch_host.path_hash(path) != checksums[path]:
-                restarts += self.full_restart_map[path]
-        services_list = list(collections.OrderedDict.fromkeys(restarts).keys())
-        for service_name in services_list:
-            self.service_stop(service_name)
-        for service_name in services_list:
-            self.service_start(service_name)
+        return ch_host.restart_on_change(
+            self.full_restart_map,
+            stopstart=True,
+            restart_functions=getattr(self, 'restart_functions', None))
 
     def restart_all(self):
         """Restart all the services configured in the self.services[]
diff --git a/unit_tests/charms_openstack/charm/test_classes.py b/unit_tests/charms_openstack/charm/test_classes.py
index 917c5a5..1603440 100644
--- a/unit_tests/charms_openstack/charm/test_classes.py
+++ b/unit_tests/charms_openstack/charm/test_classes.py
@@ -119,52 +119,13 @@ class TestOpenStackCharm(BaseOpenStackCharmTest):
         self.assertEqual(self.target.region, 'a-region')
 
     def test_restart_on_change(self):
-        from collections import OrderedDict
-        hashs = OrderedDict([
-            ('path1', 100),
-            ('path2', 200),
-            ('path3', 300),
-            ('path4', 400),
-        ])
-        self.target.restart_map = {
-            'path1': ['s1'],
-            'path2': ['s2'],
-            'path3': ['s3'],
-            'path4': ['s2', 's4'],
-        }
-        self.patch_object(chm.ch_host, 'path_hash')
-        self.path_hash.side_effect = lambda x: hashs[x]
-        self.patch_object(chm.ch_host, 'service_stop')
-        self.patch_object(chm.ch_host, 'service_start')
-        # slightly awkard, in that we need to test a context manager
-        with self.target.restart_on_change():
-            # test with no restarts
-            pass
-        self.assertEqual(self.service_stop.call_count, 0)
-        self.assertEqual(self.service_start.call_count, 0)
-
-        with self.target.restart_on_change():
-            # test with path1 and path3 restarts
-            for k in ['path1', 'path3']:
-                hashs[k] += 1
-        self.assertEqual(self.service_stop.call_count, 2)
-        self.assertEqual(self.service_start.call_count, 2)
-        self.service_stop.assert_any_call('s1')
-        self.service_stop.assert_any_call('s3')
-        self.service_start.assert_any_call('s1')
-        self.service_start.assert_any_call('s3')
-
-        # test with path2 and path4 and that s2 only gets restarted once
-        self.service_stop.reset_mock()
-        self.service_start.reset_mock()
-        with self.target.restart_on_change():
-            for k in ['path2', 'path4']:
-                hashs[k] += 1
-        self.assertEqual(self.service_stop.call_count, 2)
-        self.assertEqual(self.service_start.call_count, 2)
-        calls = [mock.call('s2'), mock.call('s4')]
-        self.service_stop.assert_has_calls(calls)
-        self.service_start.assert_has_calls(calls)
+        self.patch_object(chm.ch_host, 'restart_on_change')
+        self.restart_on_change.__enter__.return_value.name = 'a'
+        self.target.restart_on_change()
+        self.restart_on_change.assert_called_once_with(
+            {},
+            stopstart=True,
+            restart_functions=None)
 
     def test_restart_all(self):
         self.patch_object(chm.ch_host, 'service_restart')