Add context to TaskManager

Some drivers methods can require context-specific
info. Example: keystone auth enabled, auth_token
required for deploy with use vendor_passthru in API.
Possible usage:
with task_manager.acquire(context, node_id,
                          shared=False) as task:
...
    task.driver.deploy.deploy(task, node_obj)
...

def deploy(self, task, node):
...
    token = task.context.auth_token
    if token:
...

Change-Id: Id51c005499ac4013142f71703a6aa2dca067f1ad

6 files changed, 96 insertions, 55 deletions

diff --git a/ironic/tests/conductor/test_manager.py b/ironic/tests/conductor/test_manager.py
index 24ed9abcf..a5a8eb57a 100644
--- a/ironic/tests/conductor/test_manager.py
+++ b/ironic/tests/conductor/test_manager.py
@@ -81,7 +81,7 @@ class ManagerTestCase(base.DbTestCase):
         node = self.dbapi.create_node(ndict)
 
         # check that it fails if something else has locked it already
-        with task_manager.acquire(node['id'], shared=False):
+        with task_manager.acquire(self.context, node['id'], shared=False):
             node['extra'] = {'test': 'two'}
             self.assertRaises(exception.NodeLocked,
                               self.service.update_node,
@@ -214,7 +214,7 @@ class ManagerTestCase(base.DbTestCase):
         node = self.dbapi.create_node(ndict)
 
         # check if the node is locked
-        with task_manager.acquire(node['id'], shared=False):
+        with task_manager.acquire(self.context, node['id'], shared=False):
             self.assertRaises(exception.NodeLocked,
                               self.service.change_node_power_state,
                               self.context,
diff --git a/ironic/tests/conductor/test_task_manager.py b/ironic/tests/conductor/test_task_manager.py
index 92ad46c4b..2abdd4212 100644
--- a/ironic/tests/conductor/test_task_manager.py
+++ b/ironic/tests/conductor/test_task_manager.py
@@ -23,6 +23,7 @@ from testtools import matchers
 from ironic.common import exception
 from ironic.conductor import task_manager
 from ironic.db import api as dbapi
+from ironic.openstack.common import context
 from ironic.openstack.common import uuidutils
 from ironic.tests.conductor import utils as mgr_utils
 from ironic.tests.db import base
@@ -49,6 +50,7 @@ class TaskManagerTestCase(base.DbTestCase):
     def setUp(self):
         super(TaskManagerTestCase, self).setUp()
         self.dbapi = dbapi.get_instance()
+        self.context = context.get_admin_context()
         self.driver = mgr_utils.get_mocked_node_manager()
 
         self.uuids = [create_fake_node(i) for i in xrange(1, 6)]
@@ -59,7 +61,7 @@ class TaskManagerTestCase(base.DbTestCase):
 
         self.config(host='test-host')
 
-        with task_manager.acquire(uuids) as task:
+        with task_manager.acquire(self.context, uuids) as task:
             node = task.resources[0].node
             self.assertEqual(uuids[0], node.uuid)
             self.assertEqual('test-host', node.reservation)
@@ -69,7 +71,7 @@ class TaskManagerTestCase(base.DbTestCase):
 
         self.config(host='test-host')
 
-        with task_manager.acquire(uuids) as task:
+        with task_manager.acquire(self.context, uuids) as task:
             self.assertThat(task, ContainsUUIDs(uuids))
             for node in [r.node for r in task.resources]:
                 self.assertEqual('test-host', node.reservation)
@@ -78,19 +80,20 @@ class TaskManagerTestCase(base.DbTestCase):
         uuids = self.uuids[0:2]
         more_uuids = self.uuids[3:4]
 
-        with task_manager.acquire(uuids) as task:
+        with task_manager.acquire(self.context, uuids) as task:
             self.assertThat(task, ContainsUUIDs(uuids))
-            with task_manager.acquire(more_uuids) as another_task:
+            with task_manager.acquire(self.context,
+                                      more_uuids) as another_task:
                 self.assertThat(another_task, ContainsUUIDs(more_uuids))
 
     def test_get_locked_node(self):
         uuids = self.uuids[0:2]
 
         def _lock_again(u):
-            with task_manager.acquire(u):
+            with task_manager.acquire(self.context, u):
                 raise exception.IronicException("Acquired lock twice.")
 
-        with task_manager.acquire(uuids) as task:
+        with task_manager.acquire(self.context, uuids) as task:
             self.assertThat(task, ContainsUUIDs(uuids))
             self.assertRaises(exception.NodeLocked,
                               _lock_again,
@@ -100,15 +103,17 @@ class TaskManagerTestCase(base.DbTestCase):
         uuids = self.uuids[0:2]
 
         # confirm we can elevate from shared -> exclusive
-        with task_manager.acquire(uuids, shared=True) as task:
+        with task_manager.acquire(self.context, uuids, shared=True) as task:
             self.assertThat(task, ContainsUUIDs(uuids))
-            with task_manager.acquire(uuids, shared=False) as inner_task:
+            with task_manager.acquire(self.context, uuids,
+                                      shared=False) as inner_task:
                 self.assertThat(inner_task, ContainsUUIDs(uuids))
 
         # confirm someone else can still get a shared lock
-        with task_manager.acquire(uuids, shared=False) as task:
+        with task_manager.acquire(self.context, uuids, shared=False) as task:
             self.assertThat(task, ContainsUUIDs(uuids))
-            with task_manager.acquire(uuids, shared=True) as inner_task:
+            with task_manager.acquire(self.context, uuids,
+                                      shared=True) as inner_task:
                 self.assertThat(inner_task, ContainsUUIDs(uuids))
 
 
@@ -117,6 +122,7 @@ class ExclusiveLockDecoratorTestCase(base.DbTestCase):
     def setUp(self):
         super(ExclusiveLockDecoratorTestCase, self).setUp()
         self.dbapi = dbapi.get_instance()
+        self.context = context.get_admin_context()
         self.driver = mgr_utils.get_mocked_node_manager()
         self.uuids = [create_fake_node(123)]
 
@@ -127,12 +133,14 @@ class ExclusiveLockDecoratorTestCase(base.DbTestCase):
                 task.dbapi.update_node(r.node.uuid,
                                        {'power_state': 'test-state'})
 
-        with task_manager.acquire(self.uuids, shared=True) as task:
+        with task_manager.acquire(self.context, self.uuids,
+                                  shared=True) as task:
             self.assertRaises(exception.ExclusiveLockRequired,
                               do_state_change,
                               task)
 
-        with task_manager.acquire(self.uuids, shared=False) as task:
+        with task_manager.acquire(self.context, self.uuids,
+                                  shared=False) as task:
             do_state_change(task)
 
         for uuid in self.uuids:
@@ -146,12 +154,14 @@ class ExclusiveLockDecoratorTestCase(base.DbTestCase):
                                    {'power_state': 'test-state'})
 
     def test_require_exclusive_lock_on_object(self):
-        with task_manager.acquire(self.uuids, shared=True) as task:
+        with task_manager.acquire(self.context, self.uuids,
+                                  shared=True) as task:
             self.assertRaises(exception.ExclusiveLockRequired,
                               self._do_state_change,
                               task)
 
-        with task_manager.acquire(self.uuids, shared=False) as task:
+        with task_manager.acquire(self.context, self.uuids,
+                                  shared=False) as task:
             self._do_state_change(task)
 
         for uuid in self.uuids:
@@ -159,14 +169,14 @@ class ExclusiveLockDecoratorTestCase(base.DbTestCase):
             self.assertEqual('test-state', res.power_state)
 
     def test_one_node_per_task_properties(self):
-        with task_manager.acquire(self.uuids) as task:
+        with task_manager.acquire(self.context, self.uuids) as task:
             self.assertEqual(task.node, task.resources[0].node)
             self.assertEqual(task.driver, task.resources[0].driver)
             self.assertEqual(task.node_manager, task.resources[0])
 
     def test_one_node_per_task_properties_fail(self):
         self.uuids.append(create_fake_node(456))
-        with task_manager.acquire(self.uuids) as task:
+        with task_manager.acquire(self.context, self.uuids) as task:
             def get_node():
                 return task.node
 
diff --git a/ironic/tests/drivers/test_ipminative.py b/ironic/tests/drivers/test_ipminative.py
index cc6218c4f..4f5dee39a 100644
--- a/ironic/tests/drivers/test_ipminative.py
+++ b/ironic/tests/drivers/test_ipminative.py
@@ -26,6 +26,7 @@ from ironic.common import states
 from ironic.conductor import task_manager
 from ironic.db import api as db_api
 from ironic.drivers.modules import ipminative
+from ironic.openstack.common import context
 from ironic.openstack.common import jsonutils as json
 from ironic.tests import base
 from ironic.tests.conductor import utils as mgr_utils
@@ -127,6 +128,7 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
 
     def setUp(self):
         super(IPMINativeDriverTestCase, self).setUp()
+        self.context = context.get_admin_context()
         self.driver = mgr_utils.get_mocked_node_manager(
                       driver='fake_ipminative')
 
@@ -164,7 +166,8 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ipminative, '_power_on') as power_on_mock:
             power_on_mock.return_value = states.POWER_ON
 
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.driver.power.set_power_state(
                     task, self.node, states.POWER_ON)
             power_on_mock.assert_called_once_with(self.info)
@@ -173,7 +176,8 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ipminative, '_power_off') as power_off_mock:
             power_off_mock.return_value = states.POWER_OFF
 
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.driver.power.set_power_state(
                     task, self.node, states.POWER_OFF)
             power_off_mock.assert_called_once_with(self.info)
@@ -184,7 +188,8 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
             ipmicmd.set_power.return_value = {'powerstate': 'error'}
 
             self.config(retry_timeout=500, group='ipmi')
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.assertRaises(exception.PowerStateFailure,
                                   self.driver.power.set_power_state,
                                   task,
@@ -197,14 +202,15 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
             ipmicmd = ipmi_mock.return_value
             ipmicmd.set_bootdev.return_value = None
 
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.driver.power._set_boot_device(task,
                                                    self.node,
                                                    'pxe')
             ipmicmd.set_bootdev.assert_called_once_with('pxe')
 
     def test_set_boot_device_bad_device(self):
-        with task_manager.acquire([self.node['uuid']]) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']]) as task:
             self.assertRaises(exception.InvalidParameterValue,
                     self.driver.power._set_boot_device,
                     task,
@@ -215,7 +221,8 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ipminative, '_reboot') as reboot_mock:
             reboot_mock.return_value = None
 
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.driver.power.reboot(task, self.node)
             reboot_mock.assert_called_once_with(self.info)
 
@@ -225,7 +232,8 @@ class IPMINativeDriverTestCase(db_base.DbTestCase):
             ipmicmd.set_power.return_value = {'powerstate': 'error'}
 
             self.config(retry_timeout=500, group='ipmi')
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.assertRaises(exception.PowerStateFailure,
                                   self.driver.power.reboot,
                                   task,
diff --git a/ironic/tests/drivers/test_ipmitool.py b/ironic/tests/drivers/test_ipmitool.py
index 0cfa2ad0b..c06a168c8 100644
--- a/ironic/tests/drivers/test_ipmitool.py
+++ b/ironic/tests/drivers/test_ipmitool.py
@@ -25,6 +25,7 @@ import stat
 
 from oslo.config import cfg
 
+from ironic.openstack.common import context
 from ironic.openstack.common import jsonutils as json
 
 from ironic.common import exception
@@ -160,6 +161,7 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
 
     def setUp(self):
         super(IPMIToolDriverTestCase, self).setUp()
+        self.context = context.get_admin_context()
         self.dbapi = db_api.get_instance()
         self.driver = mgr_utils.get_mocked_node_manager(driver='fake_ipmitool')
 
@@ -198,7 +200,8 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
             with mock.patch.object(ipmi, '_power_off',
                                    autospec=True) as mock_off:
 
-                with task_manager.acquire([self.node['uuid']]) as task:
+                with task_manager.acquire(self.context,
+                                         [self.node['uuid']]) as task:
                     self.driver.power.set_power_state(
                             task, self.node, states.POWER_ON)
 
@@ -213,7 +216,8 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
                                    autospec=True) as mock_off:
                 mock_off.return_value = states.POWER_OFF
 
-                with task_manager.acquire([self.node['uuid']]) as task:
+                with task_manager.acquire(self.context,
+                                         [self.node['uuid']]) as task:
                     self.driver.power.set_power_state(
                             task, self.node, states.POWER_OFF)
 
@@ -228,7 +232,8 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
             with mock.patch.object(ipmi, '_power_off',
                                    autospec=True) as mock_off:
 
-                with task_manager.acquire([self.node['uuid']]) as task:
+                with task_manager.acquire(self.context,
+                                         [self.node['uuid']]) as task:
                     self.assertRaises(exception.PowerStateFailure,
                             self.driver.power.set_power_state,
                             task,
@@ -239,7 +244,7 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
                 self.assertFalse(mock_off.called)
 
     def test_set_power_invalid_state(self):
-        with task_manager.acquire([self.node['uuid']]) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']]) as task:
             self.assertRaises(exception.IronicException,
                     self.driver.power.set_power_state,
                     task,
@@ -251,13 +256,14 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
                                autospec=True) as mock_exec:
             mock_exec.return_value = [None, None]
 
-            with task_manager.acquire([self.node['uuid']]) as task:
+            with task_manager.acquire(self.context,
+                                     [self.node['uuid']]) as task:
                 self.driver.power._set_boot_device(task, self.node, 'pxe')
 
             mock_exec.assert_called_once_with(self.info, "chassis bootdev pxe")
 
     def test_set_boot_device_bad_device(self):
-        with task_manager.acquire([self.node['uuid']]) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']]) as task:
             self.assertRaises(exception.InvalidParameterValue,
                     self.driver.power._set_boot_device,
                     task,
@@ -276,7 +282,8 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
                 expected = [mock.call.power_off(self.info),
                             mock.call.power_on(self.info)]
 
-                with task_manager.acquire([self.node['uuid']]) as task:
+                with task_manager.acquire(self.context,
+                                         [self.node['uuid']]) as task:
                     self.driver.power.reboot(task, self.node)
 
                 self.assertEqual(manager.mock_calls, expected)
@@ -293,7 +300,8 @@ class IPMIToolDriverTestCase(db_base.DbTestCase):
                 expected = [mock.call.power_off(self.info),
                             mock.call.power_on(self.info)]
 
-                with task_manager.acquire([self.node['uuid']]) as task:
+                with task_manager.acquire(self.context,
+                                         [self.node['uuid']]) as task:
                     self.assertRaises(exception.PowerStateFailure,
                             self.driver.power.reboot,
                             task,
diff --git a/ironic/tests/drivers/test_pxe.py b/ironic/tests/drivers/test_pxe.py
index 4b4f30a95..8e492260d 100644
--- a/ironic/tests/drivers/test_pxe.py
+++ b/ironic/tests/drivers/test_pxe.py
@@ -425,6 +425,7 @@ class PXEDriverTestCase(db_base.DbTestCase):
 
     def setUp(self):
         super(PXEDriverTestCase, self).setUp()
+        self.context = context.get_admin_context()
         mgr_utils.get_mocked_node_manager(driver='fake_pxe')
         driver_info = INFO_DICT
         driver_info['pxe_deploy_key'] = 'fake-56789'
@@ -436,14 +437,16 @@ class PXEDriverTestCase(db_base.DbTestCase):
         self.node = self.dbapi.create_node(n)
 
     def test_validate_good(self):
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             task.resources[0].driver.deploy.validate(self.node)
 
     def test_validate_fail(self):
         info = dict(INFO_DICT)
         del info['pxe_image_source']
         self.node['driver_info'] = json.dumps(info)
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             self.assertRaises(exception.InvalidParameterValue,
                               task.resources[0].driver.deploy.validate,
                               self.node)
@@ -464,25 +467,28 @@ class PXEDriverTestCase(db_base.DbTestCase):
                     address='dd:ee:ff',
                     uuid='4fc26c0b-03f2-4d2e-ae87-c02d7f33c234',
                     node_id='123')))
-        with task_manager.acquire([self.node['uuid']]) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']]) as task:
             node_macs = pxe._get_node_mac_addresses(task, self.node)
         self.assertEqual(node_macs, ['aa:bb:cc', 'dd:ee:ff'])
 
     def test_vendor_passthru_validate_good(self):
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             task.resources[0].driver.vendor.validate(self.node,
                     method='pass_deploy_info', address='123456', iqn='aaa-bbb',
                     key='fake-56789')
 
     def test_vendor_passthru_validate_fail(self):
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             self.assertRaises(exception.InvalidParameterValue,
                               task.resources[0].driver.vendor.validate,
                               self.node, method='pass_deploy_info',
                               key='fake-56789')
 
     def test_vendor_passthru_validate_key_notmatch(self):
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             self.assertRaises(exception.InvalidParameterValue,
                               task.resources[0].driver.vendor.validate,
                               self.node, method='pass_deploy_info',
@@ -499,8 +505,8 @@ class PXEDriverTestCase(db_base.DbTestCase):
                     create_pxe_config_mock.return_value = None
                     cache_images_mock.return_value = None
 
-                    with task_manager.acquire([self.node['uuid']],
-                                              shared=False) as task:
+                    with task_manager.acquire(self.context,
+                                    [self.node['uuid']], shared=False) as task:
                         state = task.resources[0].driver.deploy.deploy(task,
                                                                     self.node)
                         get_tftp_image_info_mock.assert_called_once_with(
@@ -519,7 +525,8 @@ class PXEDriverTestCase(db_base.DbTestCase):
 
         self.useFixture(fixtures.MonkeyPatch(
                 'ironic.drivers.modules.deploy_utils.deploy', fake_deploy))
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             task.resources[0].driver.vendor.vendor_passthru(task, self.node,
                     method='pass_deploy_info', address='123456', iqn='aaa-bbb',
                     key='fake-56789')
@@ -532,7 +539,8 @@ class PXEDriverTestCase(db_base.DbTestCase):
 
         self.useFixture(fixtures.MonkeyPatch(
                 'ironic.drivers.modules.deploy_utils.deploy', fake_deploy))
-        with task_manager.acquire([self.node['uuid']], shared=True) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']],
+                                  shared=True) as task:
             self.assertRaises(exception.InstanceDeployFailure,
                             task.resources[0].driver.vendor.vendor_passthru,
                             task, self.node, method='pass_deploy_info',
@@ -606,8 +614,8 @@ class PXEDriverTestCase(db_base.DbTestCase):
                 open(deploy_kernel_path, 'w').close()
                 open(image_path, 'w').close()
 
-            with task_manager.acquire([self.node['uuid']], shared=False) \
-                    as task:
+            with task_manager.acquire(self.context, [self.node['uuid']],
+                                      shared=False) as task:
                 task.resources[0].driver.deploy.tear_down(task, self.node)
             get_tftp_image_info_mock.called_once_with(self.node)
             assert_false_path = [config_path, deploy_kernel_path, image_path,
diff --git a/ironic/tests/drivers/test_ssh.py b/ironic/tests/drivers/test_ssh.py
index 70b281950..545c3afe3 100644
--- a/ironic/tests/drivers/test_ssh.py
+++ b/ironic/tests/drivers/test_ssh.py
@@ -18,6 +18,7 @@
 import mock
 import paramiko
 
+from ironic.openstack.common import context
 from ironic.openstack.common import jsonutils as json
 
 from ironic.common import exception
@@ -391,6 +392,7 @@ class SSHDriverTestCase(db_base.DbTestCase):
 
     def setUp(self):
         super(SSHDriverTestCase, self).setUp()
+        self.context = context.get_admin_context()
         self.driver = mgr_utils.get_mocked_node_manager(driver='fake_ssh')
         self.node = db_utils.get_test_node(
                         driver='fake_ssh',
@@ -443,7 +445,7 @@ class SSHDriverTestCase(db_base.DbTestCase):
                     address='dd:ee:ff',
                     uuid='4fc26c0b-03f2-4d2e-ae87-c02d7f33c234')))
 
-        with task_manager.acquire([self.node['uuid']]) as task:
+        with task_manager.acquire(self.context, [self.node['uuid']]) as task:
             node_macs = ssh._get_nodes_mac_addresses(task, self.node)
         self.assertEqual(node_macs, ['aa:bb:cc', 'dd:ee:ff'])
 
@@ -464,8 +466,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
                     power_off_mock.return_value = None
                     power_on_mock.return_value = states.POWER_ON
 
-                    with task_manager.acquire([info['uuid']], shared=False) \
-                            as task:
+                    with task_manager.acquire(self.context, [info['uuid']],
+                                              shared=False) as task:
                         task.resources[0].driver.power.reboot(task, self.node)
 
                     self.parse_drv_info_mock.assert_called_once_with(self.node)
@@ -495,8 +497,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
                     power_off_mock.return_value = None
                     power_on_mock.return_value = states.POWER_OFF
 
-                    with task_manager.acquire([info['uuid']], shared=False) \
-                            as task:
+                    with task_manager.acquire(self.context, [info['uuid']],
+                                              shared=False) as task:
                         self.assertRaises(
                                 exception.PowerStateFailure,
                                 task.resources[0].driver.power.reboot,
@@ -521,7 +523,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
         self.get_mac_addr_mock.return_value = info['macs']
         self.get_conn_mock.return_value = self.sshclient
 
-        with task_manager.acquire([info['uuid']], shared=False) as task:
+        with task_manager.acquire(self.context, [info['uuid']],
+                                  shared=False) as task:
             self.assertRaises(
                     exception.IronicException,
                     task.resources[0].driver.power.set_power_state,
@@ -544,7 +547,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ssh, '_power_on') as power_on_mock:
             power_on_mock.return_value = states.POWER_ON
 
-            with task_manager.acquire([info['uuid']], shared=False) as task:
+            with task_manager.acquire(self.context, [info['uuid']],
+                                      shared=False) as task:
                 task.resources[0].driver.power.set_power_state(task,
                                                                self.node,
                                                                states.POWER_ON)
@@ -567,7 +571,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ssh, '_power_on') as power_on_mock:
             power_on_mock.return_value = states.POWER_OFF
 
-            with task_manager.acquire([info['uuid']], shared=False) as task:
+            with task_manager.acquire(self.context, [info['uuid']],
+                                      shared=False) as task:
                 self.assertRaises(
                         exception.PowerStateFailure,
                         task.resources[0].driver.power.set_power_state,
@@ -592,7 +597,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ssh, '_power_off') as power_off_mock:
             power_off_mock.return_value = states.POWER_OFF
 
-            with task_manager.acquire([info['uuid']], shared=False) as task:
+            with task_manager.acquire(self.context, [info['uuid']],
+                                      shared=False) as task:
                 task.resources[0].driver.power.set_power_state(task,
                         self.node, states.POWER_OFF)
 
@@ -614,7 +620,8 @@ class SSHDriverTestCase(db_base.DbTestCase):
         with mock.patch.object(ssh, '_power_off') as power_off_mock:
             power_off_mock.return_value = states.POWER_ON
 
-            with task_manager.acquire([info['uuid']], shared=False) as task:
+            with task_manager.acquire(self.context, [info['uuid']],
+                                      shared=False) as task:
                 self.assertRaises(
                         exception.PowerStateFailure,
                         task.resources[0].driver.power.set_power_state,