AMBARI-2401. Perform hostname comparison before registering the agent. (Dmitry Lysnichenko via swagle)

git-svn-id: https://svn.apache.org/repos/asf/incubator/ambari/trunk@1494299 13f79535-47bb-0310-9956-ffa450edef68

ambari-agent/conf/unix/ambari-agent

@@ -101,7 +101,7 @@ case "$1" in
           fi
         fi
         echo "Starting ambari-agent"
-        nohup $PYTHON $AGENT_SCRIPT > $OUTFILE 2>&1 &
+        nohup $PYTHON $AGENT_SCRIPT "$@" > $OUTFILE 2>&1 &
         sleep 2
         PID=$!
         echo "Verifying $AMBARI_AGENT process status..."

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

@@ -34,6 +34,7 @@ import AmbariConfig
 from security import CertificateManager
 from NetUtil import NetUtil
 import security
+import hostname
 
 
 logger = logging.getLogger()
@@ -119,7 +120,20 @@ def resolve_ambari_config():
   return config
 
 
-def perform_prestart_checks():
+def perform_prestart_checks(expected_hostname):
+  # Check if current hostname is equal to expected one (got from the server
+  # during bootstrap.
+  if expected_hostname is not None:
+    current_hostname = hostname.hostname()
+    if current_hostname != expected_hostname:
+      print("Determined hostname does not match expected. Please check agent "
+            "log for details")
+      msg = "Ambari agent machine hostname ({0}) does not match expected ambari " \
+            "server hostname ({1}). Aborting registration. Please check hostname, " \
+            "hostname -f and /etc/hosts file to confirm your " \
+            "hostname is setup correctly".format(current_hostname, expected_hostname)
+      logger.error(msg)
+      sys.exit(1)
   # Check if there is another instance running
   if os.path.isfile(ProcessHelper.pidfile):
     print("%s already exists, exiting" % ProcessHelper.pidfile)
@@ -166,8 +180,12 @@ def main():
   global config
   parser = OptionParser()
   parser.add_option("-v", "--verbose", dest="verbose", action="store_true", help="verbose log output", default=False)
+  parser.add_option("-e", "--expected-hostname", dest="expected_hostname", action="store",
+                    help="expected hostname of current host. If hostname differs, agent will fail", default=None)
   (options, args) = parser.parse_args()
 
+  expected_hostname = options.expected_hostname
+
   setup_logging(options.verbose)
 
   default_cfg = { 'agent' : { 'prefix' : '/home/ambari' } }
@@ -179,7 +197,7 @@ def main():
 
   # Check for ambari configuration file.
   config = resolve_ambari_config()
-  perform_prestart_checks()
+  perform_prestart_checks(expected_hostname)
   daemonize()
 
   killstaleprocesses()

ambari-agent/src/test/python/TestMain.py

@@ -32,6 +32,7 @@ import ConfigParser
 import os
 import tempfile
 from ambari_agent.Controller import Controller
+from optparse import OptionParser
 
 
 class TestMain(unittest.TestCase):
@@ -144,12 +145,22 @@ class TestMain(unittest.TestCase):
   @patch("sys.exit")
   @patch("os.path.isfile")
   @patch("os.path.isdir")
-  def test_perform_prestart_checks(self, isdir_mock, isfile_mock, exit_mock):
+  @patch("hostname.hostname")
+  def test_perform_prestart_checks(self, hostname_mock, isdir_mock, isfile_mock, exit_mock):
     main.config = AmbariConfig().getConfig()
+
+    # Check expected hostname test
+    hostname_mock.return_value = "test.hst"
+
+    main.perform_prestart_checks("another.hst")
+    self.assertTrue(exit_mock.called)
+
+    exit_mock.reset_mock()
+
     # Trying case if there is another instance running
     isfile_mock.return_value = True
     isdir_mock.return_value = True
-    main.perform_prestart_checks()
+    main.perform_prestart_checks(None)
     self.assertTrue(exit_mock.called)
 
     isfile_mock.reset_mock()
@@ -159,7 +170,7 @@ class TestMain(unittest.TestCase):
     # Trying case if agent prefix dir does not exist
     isfile_mock.return_value = False
     isdir_mock.return_value = False
-    main.perform_prestart_checks()
+    main.perform_prestart_checks(None)
     self.assertTrue(exit_mock.called)
 
     isfile_mock.reset_mock()
@@ -169,7 +180,7 @@ class TestMain(unittest.TestCase):
     # Trying normal case
     isfile_mock.return_value = False
     isdir_mock.return_value = True
-    main.perform_prestart_checks()
+    main.perform_prestart_checks(None)
     self.assertFalse(exit_mock.called)
 
 
@@ -223,10 +234,13 @@ class TestMain(unittest.TestCase):
   @patch.object(Controller, "__init__")
   @patch.object(Controller, "start")
   @patch.object(Controller, "join")
-  def test_main(self, join_mock, start_mock, Controller_init_mock, try_to_connect_mock, update_log_level_mock,
+  @patch("optparse.OptionParser.parse_args")
+  def test_main(self, parse_args_mock, join_mock, start_mock, Controller_init_mock, try_to_connect_mock, update_log_level_mock,
                 killstaleprocesses_mock, daemonize_mock, perform_prestart_checks_mock,
                 resolve_ambari_config_mock, stop_mock, bind_signal_handlers_mock, setup_logging_mock):
     Controller_init_mock.return_value = None
+    options = MagicMock()
+    parse_args_mock.return_value = (options, MagicMock)
 
     #testing call without command-line arguments
     main.main()
@@ -241,3 +255,10 @@ class TestMain(unittest.TestCase):
     self.assertTrue(update_log_level_mock.called)
     try_to_connect_mock.assert_called_once_with(ANY, -1, ANY)
     self.assertTrue(start_mock.called)
+
+    perform_prestart_checks_mock.reset_mock()
+
+    # Testing call with --expected-hostname parameter
+    options.expected_hostname = "test.hst"
+    main.main()
+    perform_prestart_checks_mock.assert_called_once_with(options.expected_hostname)

ambari-server/src/main/python/bootstrap.py

@@ -118,10 +118,7 @@ class SSH(threading.Thread):
     logFilePath = os.path.join(self.bootdir, self.host + ".log")
     self.writeLogToFile(logFilePath)
 
-    doneFilePath = os.path.join(self.bootdir, self.host + ".done")
-    self.writeDoneToFile(doneFilePath, str(sshstat.returncode))
-
-    logging.info("Setup agent done for host " + self.host + ", exitcode=" + str(sshstat.returncode))
+    logging.info("SSH command execution finished for host " + self.host + ", exitcode=" + str(sshstat.returncode))
     pass
 
   def writeLogToFile(self, logFilePath):
@@ -130,11 +127,6 @@ class SSH(threading.Thread):
     logFile.close
     pass
 
-  def writeDoneToFile(self, doneFilePath, returncode):
-    doneFile = open(doneFilePath, "w+")
-    doneFile.write(str(returncode))
-    doneFile.close()
-    pass
 pass
 
 
@@ -172,11 +164,31 @@ def get_difference(list1, list2):
 
 class PSSH:
   """Run SSH in parallel for a given list of hosts"""
-  def __init__(self, hosts, user, sshKeyFile, command, bootdir, errorMessage = None):
+  def __init__(self, hosts, user, sshKeyFile, bootdir, errorMessage = None, command=None, perHostCommands=None):
+    '''
+      Executes some command on all hosts via ssh. If command is equal for all
+      hosts, it should be passed as a "command" argument to PSSH constructor.
+      If command differs for different hosts, it should be passed as a
+      "perHostCommands" argument to PSSH constructor. "perHostCommands" is
+      expected to be a dictionary "hostname" -> "command", containing as many
+      entries as "hosts" list contains.
+    '''
+
+    # Checking arguments
+    # Had to include this check because wrong usage may be hard to notice without
+    # multinode cluster
+    if ((command is None) == (perHostCommands is None) or  # No any or both arguments are defined
+          (not isinstance(command, basestring)) and # "command" argument is not a string
+            (not isinstance(perHostCommands, dict) # "perHostCommands" argument is not a dictionary
+              or len(perHostCommands) != len(hosts))): # or does not contain commands for all hosts
+      raise "PSSH constructor received invalid parameters. Please " \
+            "read PSSH constructor docstring"
     self.hosts = hosts
     self.user = user
     self.sshKeyFile = sshKeyFile
+
     self.command = command
+    self.perHostCommands = perHostCommands
     self.bootdir = bootdir
     self.errorMessage = errorMessage
     self.ret = {}
@@ -191,7 +203,10 @@ class PSSH:
     for chunk in splitlist(self.hosts, 20):
       chunkstats = []
       for host in chunk:
-        ssh = SSH(self.user, self.sshKeyFile, host, self.command, self.bootdir, self.errorMessage)
+        if self.command is not None:
+          ssh = SSH(self.user, self.sshKeyFile, host, self.command, self.bootdir, self.errorMessage)
+        else:
+          ssh = SSH(self.user, self.sshKeyFile, host, self.perHostCommands[host], self.bootdir, self.errorMessage)
         ssh.start()
         chunkstats.append(ssh)
         pass
@@ -245,7 +260,6 @@ class PSCP:
         chunkstat.join(timeout)
         self.ret[chunkstat.getHost()] = chunkstat.getStatus()
       pass
-    
     pass
 pass    
     
@@ -382,7 +396,7 @@ class BootStrap:
       logging.info("Moving repo file...")
       targetDir = self.getRepoDir()
       command = self.getMoveRepoFileCommand(targetDir)
-      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir)
+      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, command=command)
       pssh.run()
       out = pssh.getstatus()
       # Preparing report about failed hosts
@@ -430,19 +444,21 @@ class BootStrap:
     else:
       return self.server_port    
     
-  def getRunSetupWithPasswordCommand(self):
-    return "sudo -S python /tmp/setupAgent.py " + os.environ[AMBARI_PASSPHRASE_VAR_NAME] + " " + self.ambariServer +\
+  def getRunSetupWithPasswordCommand(self, expected_hostname):
+    return "sudo -S python /tmp/setupAgent.py " + expected_hostname + " " + \
+           os.environ[AMBARI_PASSPHRASE_VAR_NAME] + " " + self.ambariServer +\
            " " + self.getAmbariVersion() + " " + self.getAmbariPort() + " < " + self.getPasswordFile()
 
-  def getRunSetupWithoutPasswordCommand(self):
-    return "sudo python /tmp/setupAgent.py " + os.environ[AMBARI_PASSPHRASE_VAR_NAME] + " " + self.ambariServer +\
-           " " + self.getAmbariVersion() + " " + self.getAmbariPort()
+  def getRunSetupWithoutPasswordCommand(self, expected_hostname):
+    return "sudo python /tmp/setupAgent.py " + expected_hostname + " " + \
+           os.environ[AMBARI_PASSPHRASE_VAR_NAME] + " " + self.ambariServer +\
+            " " + self.getAmbariVersion()  + " " + self.getAmbariPort()
 
-  def getRunSetupCommand(self):
+  def getRunSetupCommand(self, expected_hostname):
     if self.hasPassword():
-      return self.getRunSetupWithPasswordCommand()
+      return self.getRunSetupWithPasswordCommand(expected_hostname)
     else:
-      return self.getRunSetupWithoutPasswordCommand()
+      return self.getRunSetupWithoutPasswordCommand(expected_hostname)
 
   def runOsCheckScript(self):
     logging.info("Running os type check...")
@@ -450,7 +466,7 @@ class BootStrap:
            (self.getOsCheckScriptRemoteLocation(),
             self.getOsCheckScriptRemoteLocation(),  self.cluster_os_type)
 
-    pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir)
+    pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, command=command)
     pssh.run()
     out = pssh.getstatus()
 
@@ -458,7 +474,7 @@ class BootStrap:
     failed_current = get_difference(self.successive_hostlist, skip_failed_hosts(out))
     self.successive_hostlist = skip_failed_hosts(out)
     failed = get_difference(self.hostlist, self.successive_hostlist)
-    logging.info("Parallel ssh returns for setup agent. All failed hosts are: " + str(failed) +
+    logging.info("Parallel ssh returns for OS check. All failed hosts are: " + str(failed) +
                  ". Failed on last step: " + str(failed_current))
 
     #updating statuses
@@ -468,12 +484,14 @@ class BootStrap:
       retstatus = 0
     else:
       retstatus = 1
-    pass
+    return retstatus
 
   def runSetupAgent(self):
     logging.info("Running setup agent...")
-    command = self.getRunSetupCommand()
-    pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir)
+    perHostCommands = {}
+    for expected_hostname in self.successive_hostlist:
+      perHostCommands[expected_hostname] = self.getRunSetupCommand(expected_hostname)
+    pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, perHostCommands=perHostCommands)
     pssh.run()
     out = pssh.getstatus()
 
@@ -491,7 +509,7 @@ class BootStrap:
       retstatus = 0
     else:
       retstatus = 1
-    pass
+    return retstatus
 
   def createDoneFiles(self):
     """ Creates .done files for every host. These files are later read from Java code.
@@ -508,7 +526,7 @@ class BootStrap:
     try:
       """ Checking 'sudo' package on remote hosts """
       command = "rpm -qa | grep sudo"
-      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir, "Error: Sudo command is not available. Please install the sudo command.")
+      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, errorMessage="Error: Sudo command is not available. Please install the sudo command.", command=command)
       pssh.run()
       out = pssh.getstatus()
       # Preparing report about failed hosts
@@ -552,7 +570,7 @@ class BootStrap:
       logging.info("Changing password file mode...")
       targetDir = self.getRepoDir()
       command = "chmod 600 " + self.getPasswordFile()
-      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir)
+      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, command=command)
       pssh.run()
       out = pssh.getstatus()
       # Preparing report about failed hosts
@@ -581,7 +599,7 @@ class BootStrap:
       logging.info("Changing password file mode...")
       targetDir = self.getRepoDir()
       command = "chmod 600 " + self.getPasswordFile()
-      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, command, self.bootdir)
+      pssh = PSSH(self.successive_hostlist, self.user, self.sshkeyFile, self.bootdir, command=command)
       pssh.run()
       out = pssh.getstatus()
       # Preparing report about failed hosts
@@ -610,7 +628,7 @@ class BootStrap:
       logging.info("Deleting password file...")
       targetDir = self.getRepoDir()
       command = "rm " + self.getPasswordFile()
-      pssh = PSSH(self.hostlist_to_remove_password_file, self.user, self.sshkeyFile, command, self.bootdir)
+      pssh = PSSH(self.hostlist_to_remove_password_file, self.user, self.sshkeyFile, self.bootdir, command=command)
       pssh.run()
       out = pssh.getstatus()
       # Preparing report about failed hosts
@@ -634,7 +652,7 @@ class BootStrap:
     pass
 
   def run(self):
-    """ Copy files and run commands on remote hosts """
+    """ Copyfiles and run commands on remote hosts """
     ret1 = self.copyOsCheckScript()
     logging.info("Copying os type check script finished")
     ret2 = self.runOsCheckScript()
@@ -651,7 +669,7 @@ class BootStrap:
     ret6 = self.copyNeededFiles()
     logging.info("Copying files finished")
     ret7 = self.runSetupAgent()
-    logging.info("Running ssh command finished")
+    logging.info("Setting up agent finished")
     ret8 = 0
     if self.hasPassword() and self.hostlist_to_remove_password_file is not None:
       ret8 = self.deletePasswordFile()
@@ -659,10 +677,8 @@ class BootStrap:
     retcode = max(ret1, ret2, ret3, ret4, ret5, ret6, ret7, ret8)
     self.createDoneFiles()
     return retcode
-    pass
-  pass
-  
-  
+
+
 def main(argv=None):
   scriptDir = os.path.realpath(os.path.dirname(argv[0]))
   onlyargs = argv[1:]

ambari-server/src/main/python/setupAgent.py

@@ -37,9 +37,7 @@ AMBARI_PASSPHRASE_VAR = "AMBARI_PASSPHRASE"
 def execOsCommand(osCommand):
   osStat = subprocess.Popen(osCommand, stdout=subprocess.PIPE)
   log = osStat.communicate(0)
-  ret = {}
-  ret["exitstatus"] = osStat.returncode
-  ret["log"] = log
+  ret = {"exitstatus": osStat.returncode, "log": log}
   return ret
 
 def is_suse():
@@ -70,27 +68,31 @@ def installAgent(projectVersion):
   rpmCommand = ["yum", "-y", "install", "--nogpgcheck", "ambari-agent" + projectVersion]
   return execOsCommand(rpmCommand)
 
-def configureAgent(host):
+def configureAgent(server_hostname):
   """ Configure the agent so that it has all the configs knobs properly installed """
-  osCommand = ["sed", "-i.bak", "s/hostname=localhost/hostname=" + host + "/g", "/etc/ambari-agent/conf/ambari-agent.ini"]
+  osCommand = ["sed", "-i.bak", "s/hostname=localhost/hostname=" + server_hostname + "/g", "/etc/ambari-agent/conf/ambari-agent.ini"]
   execOsCommand(osCommand)
 
   return
 
-def runAgent(passPhrase):
+def runAgent(passPhrase, expected_hostname):
   os.environ[AMBARI_PASSPHRASE_VAR] = passPhrase
-  subprocess.call("/usr/sbin/ambari-agent start", shell=True)
+  agent_retcode = subprocess.call("/usr/sbin/ambari-agent start --expected-hostname={0}".format(expected_hostname), shell=True)
   try:
 
     ret = execOsCommand(["tail", "-20", "/var/log/ambari-agent/ambari-agent.log"])
+    try:
+      print ret['log']
+    except KeyError:
+      pass
     if not 0 == ret['exitstatus']:
       return ret['exitstatus']
-    print ret['log']
 
-    return 0
+    return agent_retcode
   except (Exception), e:
     return 1
 
+
 def getOptimalVersion(initialProjectVersion):
   if initialProjectVersion == "":
     return initialProjectVersion
@@ -153,24 +155,24 @@ def main(argv=None):
   scriptDir = os.path.realpath(os.path.dirname(argv[0]))
   # Parse the input
   onlyargs = argv[1:]
-  passPhrase = onlyargs[0]
-  hostName = onlyargs[1]
+  expected_hostname = onlyargs[0]
+  passPhrase = onlyargs[1]
+  hostname = onlyargs[2]
   projectVersion = None
   server_port = 8080
-  if len(onlyargs) > 2:
-    projectVersion = onlyargs[2]
   if len(onlyargs) > 3:
-    server_port = onlyargs[3]
+    projectVersion = onlyargs[3]
+  if len(onlyargs) > 4:
+    server_port = onlyargs[4]
   try:
     server_port = int(server_port)
   except (Exception), e:
     server_port = 8080	 
       
-
   if projectVersion is None or projectVersion == "null":
     projectVersion = ""
 
-  checkServerReachability(hostName, server_port)
+  checkServerReachability(hostname, server_port)
 
   projectVersion = getOptimalVersion(projectVersion)
   if projectVersion != "":
@@ -182,8 +184,8 @@ def main(argv=None):
     installPreReq()
     installAgent(projectVersion)
 
-  configureAgent(hostName)
-  sys.exit(runAgent(passPhrase))
+  configureAgent(hostname)
+  sys.exit(runAgent(passPhrase, expected_hostname))
 
 if __name__ == '__main__':
   logging.basicConfig(level=logging.DEBUG)

ambari-server/src/test/python/TestBootstrap.py

@@ -44,7 +44,7 @@ class TestBootstrap(TestCase):
     bootstrap.HOST_BOOTSTRAP_TIMEOUT = 1
     forever_hanging_timeout = 5
     SSH.run = lambda self: time.sleep(forever_hanging_timeout)
-    pssh = PSSH(["hostname"], "root", "sshKeyFile", "command", "bootdir")
+    pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir", command="command")
     self.assertTrue(pssh.ret == {})
     starttime = time.time()
     pssh.run()
@@ -69,15 +69,12 @@ class TestBootstrap(TestCase):
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   def test_return_error_message_for_missing_sudo_package(self, communicate_method,
-                                                         SSH_writeDoneToFile_method,
                                                          SSH_writeLogToFile_method,
                                                          SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     bootstrap = BootStrap(["hostname"], "root", "sshKeyFile", "scriptDir", "bootdir", "setupAgentFile", "ambariServer", "centos6", None, "8440")
     bootstrap.statuses = {
@@ -91,7 +88,6 @@ class TestBootstrap(TestCase):
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   @patch.object(BootStrap, "deletePasswordFile")
@@ -101,12 +97,10 @@ class TestBootstrap(TestCase):
                                                                                    deletePasswordFile_method,
                                                                                    createDoneFiles_method,
                                                                                    communicate_method,
-                                                                                   SSH_writeDoneToFile_method,
                                                                                    SSH_writeLogToFile_method,
                                                                                    SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
@@ -128,7 +122,6 @@ class TestBootstrap(TestCase):
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   @patch.object(BootStrap, "deletePasswordFile")
@@ -138,12 +131,10 @@ class TestBootstrap(TestCase):
                                                                                       deletePasswordFile_method,
                                                                                       createDoneFiles_method,
                                                                                       communicate_method,
-                                                                                      SSH_writeDoneToFile_method,
                                                                                       SSH_writeLogToFile_method,
                                                                                       SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
@@ -165,7 +156,6 @@ class TestBootstrap(TestCase):
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   @patch.object(BootStrap, "getRunSetupWithPasswordCommand")
@@ -174,12 +164,10 @@ class TestBootstrap(TestCase):
                                                                     getRunSetupWithPasswordCommand_method,
                                                                     createDoneFiles_method,
                                                                     communicate_method,
-                                                                    SSH_writeDoneToFile_method,
                                                                     SSH_writeLogToFile_method,
                                                                     SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
@@ -187,14 +175,19 @@ class TestBootstrap(TestCase):
     getMoveRepoFileWithPasswordCommand_method.return_value = ""
 
     os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
+    hosts = ["hostname"]
     bootstrap = BootStrap(["hostname"], "user", "sshKeyFile", "scriptDir", "bootdir", "setupAgentFile", "ambariServer", "centos6", None, "8440", "passwordFile")
-    ret = bootstrap.run()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.copyOsCheckScript()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.copyNeededFiles()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.runSetupAgent()
     self.assertTrue(getRunSetupWithPasswordCommand_method.called)
     self.assertTrue(getMoveRepoFileWithPasswordCommand_method.called)
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   @patch.object(BootStrap, "getRunSetupWithoutPasswordCommand")
@@ -203,12 +196,10 @@ class TestBootstrap(TestCase):
                                                                       getRunSetupWithoutPasswordCommand_method,
                                                                       createDoneFiles_method,
                                                                       communicate_method,
-                                                                      SSH_writeDoneToFile_method,
                                                                       SSH_writeLogToFile_method,
                                                                       SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
@@ -216,11 +207,18 @@ class TestBootstrap(TestCase):
     getMoveRepoFileWithoutPasswordCommand_method.return_value = ""
 
     os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
+    hosts = ["hostname"]
     bootstrap = BootStrap(["hostname"], "user", "sshKeyFile", "scriptDir", "bootdir", "setupAgentFile", "ambariServer", "centos6", None, "8440")
-    ret = bootstrap.run()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.copyOsCheckScript()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.copyNeededFiles()
+    bootstrap.successive_hostlist = hosts
+    bootstrap.runSetupAgent()
     self.assertTrue(getRunSetupWithoutPasswordCommand_method.called)
     self.assertTrue(getMoveRepoFileWithoutPasswordCommand_method.called)
 
+
   @patch.object(BootStrap, "runSetupAgent")
   @patch.object(BootStrap, "copyNeededFiles")
   @patch.object(BootStrap, "checkSudoPackage")
@@ -312,47 +310,43 @@ class TestBootstrap(TestCase):
 
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   def test_run_setup_agent_command_ends_with_project_version(self, createDoneFiles_method,
                                                              communicate_method,
-                                                             SSH_writeDoneToFile_method,
                                                              SSH_writeLogToFile_method,
                                                              SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
     os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
     version = "1.1.1"
     bootstrap = BootStrap(["hostname"], "user", "sshKeyFile", "scriptDir", "bootdir", "setupAgentFile", "ambariServer", "centos6", version, "8440")
-    runSetupCommand = bootstrap.getRunSetupCommand()
+    runSetupCommand = bootstrap.getRunSetupCommand("hostname")
     self.assertTrue(runSetupCommand.endswith(version + " 8440"))
 
+
   @patch.object(SCP, "writeLogToFile")
   @patch.object(SSH, "writeLogToFile")
-  @patch.object(SSH, "writeDoneToFile")
   @patch.object(Popen, "communicate")
   @patch.object(BootStrap, "createDoneFiles")
   def test_agent_setup_command_without_project_version(self, createDoneFiles_method,
                                                        communicate_method,
-                                                       SSH_writeDoneToFile_method,
                                                        SSH_writeLogToFile_method,
                                                        SCP_writeLogToFile_method):
     SCP_writeLogToFile_method.return_value = None
     SSH_writeLogToFile_method.return_value = None
-    SSH_writeDoneToFile_method.return_value = None
     communicate_method.return_value = ("", "")
     createDoneFiles_method.return_value = None
 
     os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
     version = None
     bootstrap = BootStrap(["hostname"], "user", "sshKeyFile", "scriptDir", "bootdir", "setupAgentFile", "ambariServer", "centos6", version, "8440")
-    runSetupCommand = bootstrap.getRunSetupCommand()
-    self.assertTrue(runSetupCommand.endswith(" 8440"))
+    runSetupCommand = bootstrap.getRunSetupCommand("hostname")
+    self.assertTrue(runSetupCommand.endswith("8440"))
+
 
 
   @patch.object(BootStrap, "createDoneFiles")
@@ -421,4 +415,63 @@ class TestBootstrap(TestCase):
     self.assertTrue(c6hstr in bootstrap.successive_hostlist)
     self.assertTrue(pssh_run_method.call_count >= 2)
     self.assertTrue(pssh_getstatus_method.call_count >= 2)
-    self.assertTrue(ret == 1)
+    self.assertTrue(ret == 1)
+
+
+  def test_PSSH_constructor_argument_validation(self):
+    dummy_command = "command"
+    dummy_dict = {
+      'hostname1' : 'c1',
+      'hostname2' : 'c2',
+      'hostname3' : 'c3'
+    }
+
+    # No any command arguments defined
+    try:
+      pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir")
+      self.fail("Should raise exception")
+    except Exception, err:
+      # Expected
+      pass
+
+    # Both command arguments defined
+    try:
+      pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir", command = dummy_command, perHostCommands = dummy_dict)
+      self.fail("Should raise exception")
+    except Exception, err:
+      # Expected
+      pass
+
+    # Invalid arguments: command dictionary has commands not for all hosts
+    inv_dict = dict(dummy_dict)
+    del inv_dict["hostname1"]
+    try:
+      pssh = PSSH(["hostname1", "hostname2", "hostname3"], "root", "sshKeyFile", "bootdir", perHostCommands=inv_dict)
+      self.fail("Should raise exception")
+    except Exception, err:
+      # Expected
+      pass
+
+    # Invalid arguments:  command dictionary instead of command
+    try:
+      pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir", command = dummy_dict)
+      self.fail("Should raise exception")
+    except Exception, err:
+      # Expected
+      pass
+
+    # Invalid arguments: single command instead of command dictionary
+    try:
+      pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir", perHostCommands = dummy_command)
+      self.fail("Should raise exception")
+    except Exception, err:
+      # Expected
+      pass
+
+    # Valid arguments: command passed
+    pssh = PSSH(["hostname"], "root", "sshKeyFile", "bootdir", command = dummy_command)
+
+    # Valid arguments: command dictionary passed
+    pssh = PSSH(["hostname1", "hostname2", "hostname3"], "root", "sshKeyFile", "bootdir", perHostCommands=dummy_dict)
+
+

ambari-server/src/test/python/TestSetupAgent.py

@@ -68,6 +68,33 @@ class TestSetupAgent(TestCase):
     pass
 
 
+  @patch.object(setup_agent, 'execOsCommand')
+  def test_configureAgent(self, execOsCommand_mock):
+    # Test if expected_hostname is passed
+    hostname = "test.hst"
+    setup_agent.configureAgent(hostname)
+    cmdStr = str(execOsCommand_mock.call_args_list[0][0])
+    self.assertTrue(hostname in cmdStr)
+
+
+  @patch.object(setup_agent, 'execOsCommand')
+  @patch("os.environ")
+  @patch("subprocess.call")
+  def test_runAgent(self, call_mock, environ_mock, execOsCommand_mock):
+    expected_hostname = "test.hst"
+    passphrase = "passphrase"
+    call_mock.return_value = 0
+    execOsCommand_mock.return_value = {'log': 'log', 'exitstatus': 0}
+
+    # Test if expected_hostname is passed
+    ret = setup_agent.runAgent(passphrase, expected_hostname)
+    cmdStr = str(call_mock.call_args_list[0][0])
+    self.assertTrue(expected_hostname in cmdStr)
+    self.assertEqual(ret, 0)
+
+
+
+
   @patch.object(setup_agent, 'is_suse')
   @patch.object(setup_agent, 'checkAgentPackageAvailabilitySuse')
   @patch.object(setup_agent, 'checkAgentPackageAvailability')