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

#!/usr/bin/env python

'''
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
'''
import StringIO
import sys
import unittest
import logging
import signal
import os
import socket
import tempfile
import platform
import ConfigParser
from ambari_commons import OSCheck
from only_for_platform import only_for_platform, get_platform, PLATFORM_WINDOWS, PLATFORM_LINUX
from mock.mock import MagicMock, patch, ANY, Mock

if get_platform() != PLATFORM_WINDOWS:
  os_distro_value = ('Suse','11','Final')
else:
  os_distro_value = ('win2012serverr2','6.3','WindowsServer')

with patch.object(OSCheck, "os_distribution", new = MagicMock(return_value = os_distro_value)):
  from ambari_agent import NetUtil, security
  from ambari_agent import ProcessHelper, main
  from ambari_agent.AmbariConfig import AmbariConfig
  from ambari_agent.PingPortListener import PingPortListener
  from ambari_agent.Controller import Controller
  from ambari_agent.DataCleaner import DataCleaner
  import ambari_agent.HeartbeatHandlers as HeartbeatHandlers
  from ambari_commons.os_check import OSConst, OSCheck

  from ambari_commons.shell import shellRunner

class TestMain(unittest.TestCase):

  def setUp(self):
    # disable stdout
    out = StringIO.StringIO()
    sys.stdout = out


  def tearDown(self):
    # enable stdout
    sys.stdout = sys.__stdout__

  @only_for_platform(PLATFORM_LINUX)
  @patch("ambari_agent.HeartbeatHandlers.HeartbeatStopHandlersLinux")
  @patch("os._exit")
  @patch("os.getpid")
  @patch.object(ProcessHelper, "stopAgent")
  def test_signal_handler(self, stopAgent_mock, os_getpid_mock, os_exit_mock, heartbeat_handler_mock):
    # testing exit of children
    main.agentPid = 4444
    os_getpid_mock.return_value = 5555
    HeartbeatHandlers.signal_handler("signum", "frame")
    heartbeat_handler_mock.set_stop.assert_called()
    os_exit_mock.reset_mock()

    # testing exit of main process
    os_getpid_mock.return_value = main.agentPid
    HeartbeatHandlers.signal_handler("signum", "frame")
    heartbeat_handler_mock.set_stop.assert_called()


  @patch.object(main.logger, "addHandler")
  @patch.object(main.logger, "setLevel")
  @patch("logging.handlers.RotatingFileHandler")
  def test_setup_logging(self, rfh_mock, setLevel_mock, addHandler_mock):
    # Testing silent mode
    main.setup_logging(False)
    self.assertTrue(addHandler_mock.called)
    setLevel_mock.assert_called_with(logging.INFO)

    addHandler_mock.reset_mock()
    setLevel_mock.reset_mock()

    # Testing verbose mode
    main.setup_logging(True)
    self.assertTrue(addHandler_mock.called)
    setLevel_mock.assert_called_with(logging.DEBUG)


  @patch.object(OSCheck, "os_distribution", new = MagicMock(return_value = os_distro_value))
  @patch.object(main.logger, "setLevel")
  @patch("logging.basicConfig")
  def test_update_log_level(self, basicConfig_mock, setLevel_mock):
    config = AmbariConfig().getConfig()

    # Testing with default setup (config file does not contain loglevel entry)
    # Log level should not be changed
    config.set('agent', 'loglevel', None)
    main.update_log_level(config)
    self.assertFalse(setLevel_mock.called)

    setLevel_mock.reset_mock()

    # Testing debug mode
    config.set('agent', 'loglevel', 'DEBUG')
    main.update_log_level(config)
    setLevel_mock.assert_called_with(logging.DEBUG)
    setLevel_mock.reset_mock()

    # Testing any other mode
    config.set('agent', 'loglevel', 'INFO')
    main.update_log_level(config)
    setLevel_mock.assert_called_with(logging.INFO)

    setLevel_mock.reset_mock()

    config.set('agent', 'loglevel', 'WRONG')
    main.update_log_level(config)
    setLevel_mock.assert_called_with(logging.INFO)

  @only_for_platform(PLATFORM_LINUX)
  @patch("signal.signal")
  def test_bind_signal_handlers(self, signal_mock):
    main.bind_signal_handlers(os.getpid())
    # Check if on SIGINT/SIGTERM agent is configured to terminate
    signal_mock.assert_any_call(signal.SIGINT, HeartbeatHandlers.signal_handler)
    signal_mock.assert_any_call(signal.SIGTERM, HeartbeatHandlers.signal_handler)
    # Check if on SIGUSR1 agent is configured to fall into debug
    signal_mock.assert_any_call(signal.SIGUSR1, HeartbeatHandlers.debug)


  @patch("os.path.exists")
  @patch("ConfigParser.RawConfigParser.read")
  def test_resolve_ambari_config(self, read_mock, exists_mock):
    # Trying case if conf file exists
    exists_mock.return_value = True
    main.resolve_ambari_config()
    self.assertTrue(read_mock.called)

    exists_mock.reset_mock()
    read_mock.reset_mock()

    # Trying case if conf file does not exist
    exists_mock.return_value = False
    main.resolve_ambari_config()
    self.assertFalse(read_mock.called)


  @patch.object(OSCheck, "os_distribution", new = MagicMock(return_value = os_distro_value))
  @patch("sys.exit")
  @patch("os.path.isfile")
  @patch("os.path.isdir")
  @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()

    if OSCheck.get_os_family() != OSConst.WINSRV_FAMILY:
      # Trying case if there is another instance running, only valid for linux
      isfile_mock.return_value = True
      isdir_mock.return_value = True
      main.perform_prestart_checks(None)
      self.assertTrue(exit_mock.called)

    isfile_mock.reset_mock()
    isdir_mock.reset_mock()
    exit_mock.reset_mock()

    # Trying case if agent prefix dir does not exist
    isfile_mock.return_value = False
    isdir_mock.return_value = False
    main.perform_prestart_checks(None)
    self.assertTrue(exit_mock.called)

    isfile_mock.reset_mock()
    isdir_mock.reset_mock()
    exit_mock.reset_mock()

    # Trying normal case
    isfile_mock.return_value = False
    isdir_mock.return_value = True
    main.perform_prestart_checks(None)
    self.assertFalse(exit_mock.called)

  @only_for_platform(PLATFORM_LINUX)
  @patch("time.sleep")
  @patch.object(shellRunner,"run")
  @patch("os._exit")
  @patch("os.path.exists")
  def test_daemonize_and_stop(self, exists_mock, _exit_mock, kill_mock, sleep_mock):
    oldpid = ProcessHelper.pidfile
    pid = str(os.getpid())
    _, tmpoutfile = tempfile.mkstemp()
    ProcessHelper.pidfile = tmpoutfile

    # Test daemonization
    main.daemonize()
    saved = open(ProcessHelper.pidfile, 'r').read()
    self.assertEqual(pid, saved)

    # Reuse pid file when testing agent stop
    # Testing normal exit
    exists_mock.return_value = False
    main.stop_agent()
    kill_mock.assert_called_with(['ambari-sudo.sh', 'kill', '-15', pid])
    _exit_mock.assert_called_with(0)

    # Restore
    kill_mock.reset_mock()
    _exit_mock.reset_mock()
    kill_mock.return_value = {'exitCode': 0, 'output': 'out', 'error': 'err'}

    # Testing exit when failed to remove pid file
    exists_mock.return_value = True
    main.stop_agent()
    kill_mock.assert_any_call(['ambari-sudo.sh', 'kill', '-15', pid])
    kill_mock.assert_any_call(['ambari-sudo.sh', 'kill', '-9', pid])
    _exit_mock.assert_called_with(1)

    # Restore
    ProcessHelper.pidfile = oldpid
    os.remove(tmpoutfile)

  @patch("os.rmdir")
  @patch("os.path.join")
  @patch('__builtin__.open')
  @patch.object(ConfigParser, "ConfigParser")
  @patch("os._exit")
  @patch("os.walk")
  @patch("os.remove")
  def test_reset(self, os_remove_mock, os_walk_mock, os_exit_mock, config_parser_mock, open_mock, os_path_join_mock, os_rmdir_mock):
    # Agent config update
    config_mock = MagicMock()
    os_walk_mock.return_value = [('/', ('',), ('file1.txt', 'file2.txt'))]
    config_parser_mock.return_value= config_mock
    config_mock.get('server', 'hostname').return_value = "old_host"
    main.reset_agent(["test", "reset", "new_hostname"])
    self.assertEqual(config_mock.get.call_count, 3)
    self.assertEqual(config_mock.set.call_count, 1)
    self.assertEqual(os_remove_mock.call_count, 2)

  @patch("os.rmdir")
  @patch("os.path.join")
  @patch('__builtin__.open')
  @patch.object(ConfigParser, "ConfigParser")
  @patch("os._exit")
  @patch("os.walk")
  @patch("os.remove")
  def test_reset_invalid_path(self, os_remove_mock, os_walk_mock, os_exit_mock, config_parser_mock, open_mock, os_path_join_mock, os_rmdir_mock):
    # Agent config file cannot be accessed
    config_mock = MagicMock()
    os_walk_mock.return_value = [('/', ('',), ('file1.txt', 'file2.txt'))]
    config_parser_mock.return_value= config_mock
    config_mock.get('server', 'hostname').return_value = "old_host"
    open_mock.side_effect = Exception("Invalid Path!")
    try:
      main.reset_agent(["test", "reset", "new_hostname"])
      self.fail("Should have thrown exception!")
    except:
      self.assertTrue(True)


  @patch.object(OSCheck, "os_distribution", new = MagicMock(return_value = os_distro_value))
  @patch.object(socket, "gethostbyname")
  @patch.object(main, "setup_logging")
  @patch.object(main, "bind_signal_handlers")
  @patch.object(main, "stop_agent")
  @patch.object(AmbariConfig, "getConfigFile")
  @patch.object(main, "perform_prestart_checks")
  @patch.object(main, "daemonize")
  @patch.object(main, "update_log_level")
  @patch.object(NetUtil.NetUtil, "try_to_connect")
  @patch.object(Controller, "__init__")
  @patch.object(Controller, "start")
  @patch.object(Controller, "join")
  @patch("optparse.OptionParser.parse_args")
  @patch.object(DataCleaner,"start")
  @patch.object(DataCleaner,"__init__")
  @patch.object(PingPortListener,"start")
  @patch.object(PingPortListener,"__init__")
  def test_main(self, ping_port_init_mock, ping_port_start_mock, data_clean_init_mock,data_clean_start_mock,
                parse_args_mock, join_mock, start_mock, Controller_init_mock, try_to_connect_mock,
                update_log_level_mock, daemonize_mock, perform_prestart_checks_mock,
                ambari_config_mock,
                stop_mock, bind_signal_handlers_mock,
                setup_logging_mock, socket_mock):
    data_clean_init_mock.return_value = None
    Controller_init_mock.return_value = None
    ping_port_init_mock.return_value = None
    options = MagicMock()
    parse_args_mock.return_value = (options, MagicMock)
    try_to_connect_mock.return_value = (0, True)
    # use default unix config
    ambari_config_mock.return_value = os.path.abspath("../../../conf/unix/ambari-agent.ini")
    #testing call without command-line arguments

    main.main()

    self.assertTrue(setup_logging_mock.called)
    self.assertTrue(bind_signal_handlers_mock.called)
    if OSCheck.get_os_family() != OSConst.WINSRV_FAMILY:
      self.assertTrue(stop_mock.called)
    #self.assertTrue(resolve_ambari_config_mock.called)
    self.assertTrue(perform_prestart_checks_mock.called)
    if OSCheck.get_os_family() != OSConst.WINSRV_FAMILY:
      self.assertTrue(daemonize_mock.called)
    self.assertTrue(update_log_level_mock.called)
    try_to_connect_mock.assert_called_once_with(ANY, -1, ANY)
    self.assertTrue(start_mock.called)
    self.assertTrue(data_clean_init_mock.called)
    self.assertTrue(data_clean_start_mock.called)
    self.assertTrue(ping_port_init_mock.called)
    self.assertTrue(ping_port_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)