AMBARI-4369 Restart command does not clear stale configs flag (dsen)

ambari-agent/src/main/python/ambari_agent/ActionQueue.py

@@ -51,6 +51,8 @@ class ActionQueue(threading.Thread):
   ROLE_COMMAND_INSTALL = 'INSTALL'
   ROLE_COMMAND_START = 'START'
   ROLE_COMMAND_STOP = 'STOP'
+  ROLE_COMMAND_CUSTOM_COMMAND = 'CUSTOM_COMMAND'
+  CUSTOM_COMMAND_RESTART = 'RESTART'
 
   IN_PROGRESS_STATUS = 'IN_PROGRESS'
   COMPLETED_STATUS = 'COMPLETED'
@@ -194,7 +196,10 @@ class ActionQueue(threading.Thread):
       if command.has_key('roleCommand') and \
         (command['roleCommand'] == self.ROLE_COMMAND_START or \
         (command['roleCommand'] == self.ROLE_COMMAND_INSTALL \
-        and component in LiveStatus.CLIENT_COMPONENTS)):
+        and component in LiveStatus.CLIENT_COMPONENTS) or \
+        (command['roleCommand'] == self.ROLE_COMMAND_CUSTOM_COMMAND and \
+        command['hostLevelParams'].has_key('custom_command') and \
+        command['hostLevelParams']['custom_command'] == self.CUSTOM_COMMAND_RESTART)):
         configHandler.copy_to_component(command['role'])
         roleResult['configurationTags'] = configHandler.read_actual_component(command['role'])
     self.commandStatuses.put_command_status(command, roleResult)

ambari-agent/src/test/python/ambari_agent/TestActionQueue.py

@@ -32,6 +32,9 @@ from threading import Thread
 from mock.mock import patch, MagicMock, call
 from ambari_agent.StackVersionsFileHandler import StackVersionsFileHandler
 from ambari_agent.CustomServiceOrchestrator import CustomServiceOrchestrator
+from ambari_agent.PythonExecutor import PythonExecutor
+from ambari_agent.CommandStatusDict import CommandStatusDict
+from ambari_agent.ActualConfigHandler import ActualConfigHandler
 
 
 class TestActionQueue(TestCase):
@@ -125,6 +128,18 @@ class TestActionQueue(TestCase):
     'configurations':{}
   }
 
+  datanode_restart_command = {
+    'commandType': 'EXECUTION_COMMAND',
+    'role': u'DATANODE',
+    'roleCommand': u'CUSTOM_COMMAND',
+    'commandId': '1-1',
+    'taskId': 9,
+    'clusterName': u'cc',
+    'serviceName': u'HDFS',
+    'configurations':{'global' : {}},
+    'configurationTags':{'global' : { 'tag': 'v123' }},
+    'hostLevelParams':{'custom_command': 'RESTART'}
+  }
 
   @patch.object(ActionQueue, "process_command")
   @patch.object(Queue, "get")
@@ -353,6 +368,46 @@ class TestActionQueue(TestCase):
     self.assertEqual(len(report['reports']), 0)
 
 
+  @patch.object(CustomServiceOrchestrator, "runCommand")
+  @patch("CommandStatusDict.CommandStatusDict")
+  @patch.object(ActionQueue, "status_update_callback")
+  @patch.object(ActionQueue, "determine_command_format_version")
+  def test_store_configuration_tags(self, determine_command_format_version_mock,
+                                    status_update_callback_mock,
+                                    command_status_dict_mock,
+                                    cso_runCommand_mock):
+    determine_command_format_version_mock.return_value = 2
+    custom_service_orchestrator_execution_result_dict = {
+      'stdout': 'out',
+      'stderr': 'stderr',
+      'structuredOut' : '',
+      'exitcode' : 0
+    }
+    cso_runCommand_mock.return_value = custom_service_orchestrator_execution_result_dict
+
+    config = AmbariConfig().getConfig()
+    tempdir = tempfile.gettempdir()
+    config.set('agent', 'prefix', tempdir)
+    actionQueue = ActionQueue(config, 'dummy_controller')
+    actionQueue.execute_command(self.datanode_restart_command)
+    report = actionQueue.result()
+    expected = {'actionId': '1-1',
+                'clusterName': u'cc',
+                'configurationTags': {'global' : { 'tag': 'v123' }},
+                'exitCode': 0,
+                'role': u'DATANODE',
+                'roleCommand': u'CUSTOM_COMMAND',
+                'serviceName': u'HDFS',
+                'status': 'COMPLETED',
+                'stderr': 'stderr',
+                'stdout': 'out',
+                'structuredOut': '',
+                'taskId': 9
+    }
+    # Agent caches configurationTags if custom_command RESTART completed
+    self.assertEqual(len(report['reports']), 1)
+    self.assertEqual(expected, report['reports'][0])
+
   @patch.object(ActionQueue, "status_update_callback")
   @patch.object(ActionQueue, "determine_command_format_version")
   @patch.object(StackVersionsFileHandler, "read_stack_version")