slider-agent/src/test/python/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 NetUtil, security
from mock.mock import MagicMock, patch, ANY
import unittest
import ProcessHelper, main
import logging
import signal
from AgentConfig import AgentConfig
import ConfigParser
import os
import tempfile
from Controller import Controller
from Registry import Registry
from optparse import OptionParser
import platform

IS_WINDOWS = platform.system() == "Windows"

logger = logging.getLogger()

class TestMain(unittest.TestCase):
  def setUp(self):
    # disable stdout
    out = StringIO.StringIO()
    sys.stdout = out


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

  @patch("os._exit")
  @patch("os.getpid")
  @patch.object(ProcessHelper, "stopAgent")
  def test_signal_handler(self, stopAgent_mock, os_getpid_mock, os_exit_mock):
    # testing exit of children
    main.agentPid = 4444
    os_getpid_mock.return_value = 5555
    main.signal_handler("signum", "frame")
    self.assertTrue(os_exit_mock.called)

    os_exit_mock.reset_mock()

    # testing exit of main process
    os_getpid_mock.return_value = main.agentPid
    main.signal_handler("signum", "frame")
    self.assertFalse(os_exit_mock.called)
    self.assertTrue(stopAgent_mock.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
    _, tmpoutfile = tempfile.mkstemp()
    main.setup_logging(False, tmpoutfile)
    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, tmpoutfile)
    self.assertTrue(addHandler_mock.called)
    setLevel_mock.assert_called_with(logging.DEBUG)


  @patch.object(main.logger, "setLevel")
  @patch("logging.basicConfig")
  def test_update_log_level(self, basicConfig_mock, setLevel_mock):
    config = AgentConfig("", "")
    _, tmpoutfile = tempfile.mkstemp()

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

    setLevel_mock.reset_mock()

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

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

    setLevel_mock.reset_mock()

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

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


  @patch("os.path.exists")
  @patch("ConfigParser.RawConfigParser.read")
  def test_resolve_config(self, read_mock, exists_mock):
    config = AgentConfig("", "")
    # Trying case if conf file exists
    exists_mock.return_value = True
    main.update_config_from_file(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.update_config_from_file(config)
    self.assertFalse(read_mock.called)


  @patch("os.remove")
  @patch("sys.exit")
  @patch("os.path.isfile")
  @patch("os.path.isdir")
  def test_perform_prestart_checks(self, isdir_mock, isfile_mock,
                                   exit_mock, remove_mock):
    main.config = AgentConfig("", "")

    # Trying case if there is another instance running
    isfile_mock.return_value = True
    isdir_mock.return_value = True
    main.perform_prestart_checks(main.config)
    self.assertFalse(exit_mock.called)
    self.assertTrue(remove_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(main.config)
    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(main.config)
    self.assertFalse(exit_mock.called)

  if not IS_WINDOWS:
    @patch("time.sleep")
    @patch("os.kill")
    @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.write_pid()
      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(int(pid), signal.SIGTERM)

      # Restore
      kill_mock.reset_mock()
      _exit_mock.reset_mock()

      # Testing exit when failed to remove pid file
      exists_mock.return_value = True
      main.stop_agent()
      kill_mock.assert_any_call(int(pid), signal.SIGTERM)
      kill_mock.assert_any_call(int(pid), signal.SIGKILL)
      _exit_mock.assert_called_with(1)

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

  @patch.object(Registry, "readAMHostPort")
  @patch.object(main, "setup_logging")
  @patch.object(main, "bind_signal_handlers")
  @patch.object(main, "update_config_from_file")
  @patch.object(main, "perform_prestart_checks")
  @patch.object(main, "write_pid")
  @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(Controller, "is_alive")
  def test_main(self, isAlive_mock, parse_args_mock, join_mock, start_mock,
                Controller_init_mock, try_to_connect_mock,
                update_log_level_mock, write_pid_mock,
                perform_prestart_checks_mock,
                update_config_from_file_mock,
                bind_signal_handlers_mock, setup_logging_mock,
                readAMHostPort_mock):
    Controller_init_mock.return_value = None
    isAlive_mock.return_value = False
    options = MagicMock()
    parse_args_mock.return_value = (options, MagicMock)
    readAMHostPort_mock.return_value = ("host1", 101, 100)

    tmpdir = tempfile.gettempdir()

    #testing call without command-line arguments
    os.environ["AGENT_WORK_ROOT"] = os.path.join(tmpdir, "work")
    os.environ["AGENT_LOG_ROOT"] = ",".join([os.path.join(tmpdir, "log"),os.path.join(tmpdir, "log2")])
    try_to_connect_mock.return_value = 1
    main.main()

    self.assertTrue(setup_logging_mock.called)
    if not IS_WINDOWS:
      self.assertTrue(bind_signal_handlers_mock.called)
    self.assertTrue(update_config_from_file_mock.called)
    self.assertTrue(perform_prestart_checks_mock.called)
    self.assertTrue(write_pid_mock.called)
    self.assertTrue(update_log_level_mock.called)
    self.assertTrue(options.log_folder == os.path.join(tmpdir, "log"))
    try_to_connect_mock.assert_called_once_with('https://host1:101/ws/v1/slider/agents/', 3, ANY)
    self.assertTrue(start_mock.called)

  class AgentOptions:
      def __init__(self, label, zk_quorum, zk_reg_path, verbose, debug):
          self.label = label
          self.zk_quorum = zk_quorum
          self.zk_reg_path = zk_reg_path
          self.verbose = verbose
          self.debug = debug

  @patch.object(Registry, "readAMHostPort")
  @patch("time.sleep")
  @patch.object(main, "setup_logging")
  @patch.object(main, "bind_signal_handlers")
  @patch.object(main, "stop_agent")
  @patch.object(main, "update_config_from_file")
  @patch.object(main, "perform_prestart_checks")
  @patch.object(main, "write_pid")
  @patch.object(main, "update_log_level")
  @patch.object(NetUtil.NetUtil, "try_to_connect")
  @patch.object(AgentConfig, 'set')
  @patch.object(Controller, "__init__")
  @patch.object(Controller, "start")
  @patch.object(Controller, "join")
  @patch.object(Controller, "is_alive")
  @patch("optparse.OptionParser.parse_args")
  def test_main2(self, parse_args_mock, isAlive_mock, join_mock, start_mock,
                Controller_init_mock, AgentConfig_set_mock,
                try_to_connect_mock,
                update_log_level_mock, write_pid_mock,
                perform_prestart_checks_mock,
                update_config_from_file_mock, stop_mock,
                bind_signal_handlers_mock, setup_logging_mock,
                time_sleep_mock, readAMHostPort_mock):
      Controller_init_mock.return_value = None
      isAlive_mock.return_value = False
      parse_args_mock.return_value = (
          TestMain.AgentOptions("agent", "host1:2181", "/registry/org-apache-slider/cl1", True, ""), [])
      tmpdir = tempfile.gettempdir()
      time_sleep_mock.return_value = None
      readAMHostPort_mock.return_value = (None, None, None)

      #testing call without command-line arguments
      os.environ["AGENT_WORK_ROOT"] = os.path.join(tmpdir, "work")
      os.environ["AGENT_LOG_ROOT"] = os.path.join(tmpdir, "log")
      main.main()
      self.assertTrue(AgentConfig_set_mock.call_count == 3)
      self.assertTrue(readAMHostPort_mock.call_count == 10)
      AgentConfig_set_mock.assert_any_call("server", "zk_quorum", "host1:2181")
      AgentConfig_set_mock.assert_any_call("server", "zk_reg_path", "/registry/org-apache-slider/cl1")


  def test_config1(self):
    config = AgentConfig("", "")
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 5)
    self.assertEqual(window, 5)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, '')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 0)
    self.assertEqual(window, 0)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, '33')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 0)
    self.assertEqual(window, 0)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, '-4,-6')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 0)
    self.assertEqual(window, 0)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, 'wd,er')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 0)
    self.assertEqual(window, 0)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, '2,20')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 2)
    self.assertEqual(window, 20)

    config.set(AgentConfig.COMMAND_SECTION, AgentConfig.AUTO_RESTART, ' 2, 30')
    (max, window) = config.getErrorWindow()
    self.assertEqual(max, 0)
    self.assertEqual(window, 0)


if __name__ == "__main__":
  logging.basicConfig(format='%(asctime)s %(message)s', level=logging.DEBUG)
  unittest.main()