Python MagicMock.assert_called_with方法代码示例

# 需要导入模块: from mock.mock import MagicMock [as 别名]
# 或者: from mock.mock.MagicMock import assert_called_with [as 别名]
  def test_recoveryHbCmd(self, dumpsMock, sleepMock, event_mock, exit_mock):

    out = StringIO.StringIO()
    sys.stdout = out

    hearbeat = MagicMock()
    self.controller.heartbeat = hearbeat
    event_mock.return_value = False
    dumpsMock.return_value = "data"

    sendRequest = MagicMock(name="sendRequest")
    self.controller.sendRequest = sendRequest
    addToQueue = MagicMock(name="addToQueue")
    addToStatusQueue = MagicMock(name="addToStatusQueue")
    self.addToQueue = addToQueue
    self.addToStatusQueue = addToStatusQueue

    process_execution_commands = MagicMock(name="process_execution_commands")
    self.controller.recovery_manager.process_execution_commands = process_execution_commands
    process_status_commands = MagicMock(name="process_status_commands")
    self.controller.recovery_manager.process_status_commands = process_status_commands
    set_paused = MagicMock(name = "set_paused")
    self.controller.recovery_manager.set_paused = set_paused

    self.controller.responseId = 0
    response = {"responseId":1,
                "statusCommands": "commands2",
                "executionCommands" : "commands1",
                "log":"",
                "exitstatus":"0",
                "hasPendingTasks": True}
    sendRequest.return_value = response

    def one_heartbeat(*args, **kwargs):
      self.controller.DEBUG_STOP_HEARTBEATING = True
      return response

    sendRequest.side_effect = one_heartbeat

    actionQueue = MagicMock()
    actionQueue.isIdle.return_value = True

    # one successful request, after stop
    self.controller.actionQueue = actionQueue
    self.controller.heartbeatWithServer()
    self.assertTrue(sendRequest.called)
    self.assertTrue(process_execution_commands.called)
    self.assertFalse(process_status_commands.called)
    process_execution_commands.assert_called_with("commands1")
    set_paused.assert_called_with(True)

    self.controller.heartbeatWithServer()
    sys.stdout = sys.__stdout__
    self.controller.sendRequest = Controller.Controller.sendRequest
    self.controller.sendRequest = Controller.Controller.addToQueue
    self.controller.sendRequest = Controller.Controller.addToStatusQueue
    pass

# 需要导入模块: from mock.mock import MagicMock [as 别名]
# 或者: from mock.mock.MagicMock import assert_called_with [as 别名]
  def test_heartbeatWithServer(self, dumpsMock, sleepMock, event_mock, exit_mock):
    out = StringIO.StringIO()
    sys.stdout = out

    hearbeat = MagicMock()
    self.controller.heartbeat = hearbeat
    event_mock.return_value = False
    dumpsMock.return_value = "data"

    sendRequest = MagicMock(name="sendRequest")
    self.controller.sendRequest = sendRequest

    self.controller.responseId = 1
    response = {"responseId":"2", "restartAgent":"false"}
    sendRequest.return_value = response

    def one_heartbeat(*args, **kwargs):
      self.controller.DEBUG_STOP_HEARTBEATING = True
      return response

    sendRequest.side_effect = one_heartbeat

    actionQueue = MagicMock()
    actionQueue.isIdle.return_value = True

    # one successful request, after stop
    self.controller.actionQueue = actionQueue
    self.controller.alert_scheduler_handler = MagicMock()
    self.controller.heartbeatWithServer()
    self.assertTrue(sendRequest.called)

    calls = []
    def retry(*args, **kwargs):
      if len(calls) == 0:
        calls.append(1)
        response["responseId"] = "3"
        raise Exception()
      if len(calls) > 0:
        self.controller.DEBUG_STOP_HEARTBEATING = True
      return response

    # exception, retry, successful and stop
    sendRequest.side_effect = retry
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    self.assertEqual(1, self.controller.DEBUG_SUCCESSFULL_HEARTBEATS)

    # retry registration
    self.controller.responseId = 2
    response["registrationCommand"] = "true"
    sendRequest.side_effect = one_heartbeat
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    self.assertTrue(self.controller.repeatRegistration)

    # components are not mapped
    self.controller.responseId = 2
    response["registrationCommand"] = "false"
    response["hasMappedComponents"] = False
    sendRequest.side_effect = one_heartbeat
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    self.assertFalse(self.controller.hasMappedComponents)

    # components are mapped
    self.controller.responseId = 2
    response["hasMappedComponents"] = True
    sendRequest.side_effect = one_heartbeat
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    self.assertTrue(self.controller.hasMappedComponents)

    # components are mapped
    self.controller.responseId = 2
    del response["hasMappedComponents"]
    sendRequest.side_effect = one_heartbeat
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    self.assertTrue(self.controller.hasMappedComponents)

    # wrong responseId => restart
    self.controller.responseId = 2
    response = {"responseId":"2", "restartAgent":"false"}

    restartAgent = MagicMock(name="restartAgent")
    self.controller.restartAgent = restartAgent
    self.controller.DEBUG_STOP_HEARTBEATING = False
    self.controller.heartbeatWithServer()

    restartAgent.assert_called_with()

    # executionCommands
    self.controller.responseId = 1
    addToQueue = MagicMock(name="addToQueue")
    self.controller.addToQueue = addToQueue
#.........这里部分代码省略.........

# 需要导入模块: from mock.mock import MagicMock [as 别名]
# 或者: from mock.mock.MagicMock import assert_called_with [as 别名]
  def test_post_upgrade_restart(self, time_mock):
    # load the NN and JN JMX files so that the urllib2.urlopen mock has data
    # to return
    num_journalnodes = 3
    journalnode_jmx_file = os.path.join(RMFTestCase._getStackTestsFolder(),
      self.UPGRADE_STACK_VERSION, "configs", "journalnode-upgrade-jmx.json")

    namenode_jmx_file = os.path.join(RMFTestCase._getStackTestsFolder(),
      self.UPGRADE_STACK_VERSION, "configs", "journalnode-upgrade-namenode-jmx.json")

    namenode_status_active_file = os.path.join(RMFTestCase._getStackTestsFolder(),
      self.UPGRADE_STACK_VERSION, "configs", "journalnode-upgrade-namenode-status-active.json")

    namenode_status_standby_file = os.path.join(RMFTestCase._getStackTestsFolder(),
      self.UPGRADE_STACK_VERSION, "configs", "journalnode-upgrade-namenode-status-standby.json")

    journalnode_jmx = open(journalnode_jmx_file, 'r').read()
    namenode_jmx = open(namenode_jmx_file, 'r').read()
    namenode_status_active = open(namenode_status_active_file, 'r').read()
    namenode_status_standby = open(namenode_status_standby_file, 'r').read()

    import utils
    import urllib2
    from namenode_ha_state import NamenodeHAState

    url_stream_mock = MagicMock()
    url_stream_mock.read.side_effect = (num_journalnodes * [namenode_jmx, journalnode_jmx])
    urlopen_mock = MagicMock(return_value = url_stream_mock)
    #urlopen_mock.return_value = url_stream_mock
    curl_krb_request_mock = MagicMock(side_effect=(num_journalnodes * [(namenode_jmx, "", 1), (journalnode_jmx, "", 1)]))
    get_address_mock = MagicMock(return_value="c6406.ambari.apache.org")
    with patch.object(utils, "curl_krb_request", curl_krb_request_mock):
      with patch.object(urllib2, "urlopen", urlopen_mock):
       with patch.object(NamenodeHAState, "get_address", get_address_mock):
         self.executeScript(self.COMMON_SERVICES_PACKAGE_DIR + "/scripts/journalnode.py",
           classname = "JournalNode", command = "post_upgrade_restart",
           config_file = "journalnode-upgrade.json",
           checked_call_mocks = [(0, str(namenode_status_active)), (0, str(namenode_status_standby))],
           hdp_stack_version = self.UPGRADE_STACK_VERSION,
           target = RMFTestCase.TARGET_COMMON_SERVICES )

    # ensure that the mock was called with the http-style version of the URL
    urlopen_mock.assert_called
    urlopen_mock.assert_called_with("http://c6407.ambari.apache.org:8480/jmx")

    url_stream_mock.reset_mock()
    curl_krb_request_mock.reset_mock()
    get_address_mock.reset_mock()

    # now try with HDFS on SSL
    with patch.object(utils, "curl_krb_request", curl_krb_request_mock):
      with patch.object(urllib2, "urlopen", urlopen_mock):
        with patch.object(NamenodeHAState, "get_address", get_address_mock):
         self.executeScript(self.COMMON_SERVICES_PACKAGE_DIR + "/scripts/journalnode.py",
           classname = "JournalNode", command = "post_upgrade_restart",
           config_file = "journalnode-upgrade-hdfs-secure.json",
           checked_call_mocks = [(0, str(namenode_status_active)), (0, str(namenode_status_standby))],
           hdp_stack_version = self.UPGRADE_STACK_VERSION,
           target = RMFTestCase.TARGET_COMMON_SERVICES )

    # ensure that the mock was called with the http-style version of the URL
    curl_krb_request_mock.assert_called
    curl_krb_request_mock.assert_called_with("/tmp", "/etc/security/keytabs/smokeuser.headless.keytab",
                                             "[email protected]", "https://c6407.ambari.apache.org:8481/jmx",
                                             "jn_upgrade", "/usr/bin/kinit", False, None, "ambari-qa")