1 files changed, 66 insertions, 1 deletions
diff --git a/ironic/tests/unit/drivers/modules/test_agent_client.py b/ironic/tests/unit/drivers/modules/test_agent_client.py
index 1f8cc1f6a..cb794c94c 100644
--- a/ironic/tests/unit/drivers/modules/test_agent_client.py
+++ b/ironic/tests/unit/drivers/modules/test_agent_client.py
@@ -43,7 +43,8 @@ class MockResponse(object):
 
 
 class MockCommandStatus(MockResponse):
-    def __init__(self, status, name='fake', error=None):
+    def __init__(self, status, name='fake', error=None,
+                 status_code=http_client.OK):
         super().__init__({
             'commands': [
                 {'command_name': name,
@@ -54,6 +55,12 @@ class MockCommandStatus(MockResponse):
         })
 
 
+class MockFault(MockResponse):
+    def __init__(self, faultstring, status_code=http_client.BAD_REQUEST):
+        super().__init__({'faultstring': faultstring},
+                         status_code=status_code)
+
+
 class MockNode(object):
     def __init__(self):
         self.uuid = 'uuid'
@@ -73,6 +80,9 @@ class MockNode(object):
             'driver_info': self.driver_info,
         }
 
+    def save(self):
+        pass
+
 
 class TestAgentClient(base.TestCase):
     def setUp(self):
@@ -606,6 +616,61 @@ class TestAgentClientAttempts(base.TestCase):
             verify=True)
 
     @mock.patch.object(retrying.time, 'sleep', autospec=True)
+    def test__command_succeed_after_agent_token(self, mock_sleep):
+        self.config(require_agent_token=False)
+        mock_sleep.return_value = None
+        error = 'Unknown Argument: "agent_token"'
+        response_data = {'status': 'ok'}
+        method = 'standby.run_image'
+        image_info = {'image_id': 'test_image'}
+        params = {'image_info': image_info}
+        i_info = self.node.driver_internal_info
+        i_info['agent_secret_token'] = 'meowmeowmeow'
+        self.client.session.post.side_effect = [
+            MockFault(error),
+            MockResponse(response_data),
+        ]
+
+        response = self.client._command(self.node, method, params)
+        self.assertEqual(2, self.client.session.post.call_count)
+        self.assertEqual(response, response_data)
+        self.client.session.post.assert_called_with(
+            self.client._get_command_url(self.node),
+            data=self.client._get_command_body(method, params),
+            params={'wait': 'false'},
+            timeout=60,
+            verify=True)
+        self.assertNotIn('agent_secret_token', self.node.driver_internal_info)
+
+    @mock.patch.object(retrying.time, 'sleep', autospec=True)
+    def test__command_fail_agent_token_required(self, mock_sleep):
+        self.config(require_agent_token=True)
+        mock_sleep.return_value = None
+        error = 'Unknown Argument: "agent_token"'
+        method = 'standby.run_image'
+        image_info = {'image_id': 'test_image'}
+        params = {'image_info': image_info}
+        i_info = self.node.driver_internal_info
+        i_info['agent_secret_token'] = 'meowmeowmeow'
+        self.client.session.post.side_effect = [
+            MockFault(error)
+        ]
+
+        self.assertRaises(exception.AgentAPIError,
+                          self.client._command,
+                          self.node, method, params)
+        self.assertEqual(1, self.client.session.post.call_count)
+        self.client.session.post.assert_called_with(
+            self.client._get_command_url(self.node),
+            data=self.client._get_command_body(method, params),
+            params={'wait': 'false', 'agent_token': 'meowmeowmeow'},
+            timeout=60,
+            verify=True)
+        self.assertEqual(
+            'meowmeowmeow',
+            self.node.driver_internal_info.get('agent_secret_token'))
+
+    @mock.patch.object(retrying.time, 'sleep', autospec=True)
     def test__command_succeed_after_one_timeout(self, mock_sleep):
         mock_sleep.return_value = None
         error = 'Connection Timeout'