Python bladerunner.Bladerunner类代码示例

def test_run_single_error():
    """Ensure an error is passed back during run_single on connect errors."""

    runner = Bladerunner({
        "username": "dudeguy",
        "password": "hunter111",
        "progressbar": True,
        "port": 2212,
    })

    runner.progress = ProgressBar(1, runner.options)

    with patch.object(runner, "connect", return_value=(None, -3)) as p_connect:
        with patch.object(runner.progress, "update") as p_update:
            ret = runner._run_single("nowhere")

    p_connect.assert_called_once_with(
        "nowhere",
        "dudeguy",
        "hunter111",
        2212,
    )

    # if the progressbar should tick regardless of success
    assert p_update.called
    assert ret == {"name": "nowhere", "results": [("login", runner.errors[2])]}

def test_run_interactive_returns():
    """Confirm the dict return when print_results is False."""

    runner = Bladerunner({"threads": 17})
    fake_result = Mock()
    fake_result.result = Mock(return_value="some result")

    mock_con = Mock()
    runner.interactive_hosts = {"host": mock_con}

    fake_thread = Mock()
    fake_thread.__enter__ = fake_context
    fake_thread.__exit__ = fake_context
    fake_thread.submit = Mock(return_value=fake_result)

    threadpool_mock = patch.object(
        base,
        "ThreadPoolExecutor",
        return_value=fake_thread,
    )

    with patch.object(runner, "setup_interactive") as mock_setup:
        with threadpool_mock as mock_threadpool:
            results = runner.run_interactive("ok", "host", print_results=False)

    mock_setup.assert_called_once_with("host")
    mock_threadpool.assert_called_once_with(max_workers=17)

    assert results == {"host": "some result"}

def test_interactive_function_success():
    """Ensure the function is called when it passes inspection."""

    def good_function(session):
        return session.run("who")

    runner = Bladerunner()
    mock_hosts = Mock()
    # we need to mock out the interactive_hosts dict b/c dict.values will
    # return a new object every time, making the assert_called_once_with fail
    runner.interactive_hosts = mock_hosts

    map_mock = patch.object(
        base.ThreadPoolExecutor,
        "map",
        return_value=["totes fake"],
    )

    with patch.object(runner, "setup_interactive") as mock_setup:
        with map_mock as mock_map:
            res = runner.run_interactive_function(good_function, "some host")

    mock_setup.assert_called_once_with("some host")
    assert mock_map.call_count == 1
    mock_map.assert_called_once_with(
        good_function,
        mock_hosts.values(),
    )

    assert res == ["totes fake"]

def test_run_interactive(capfd):
    """Ensure the calls to run a command on a list of hosts interactively."""

    runner = Bladerunner({"threads": 14})
    fake_result = Mock()
    fake_result.result = Mock(return_value="fake results")

    mock_con = Mock()
    runner.interactive_hosts = {"fake host": mock_con}

    fake_thread = Mock()
    fake_thread.__enter__ = fake_context
    fake_thread.__exit__ = fake_context
    fake_thread.submit = Mock(return_value=fake_result)

    threadpool_mock = patch.object(
        base,
        "ThreadPoolExecutor",
        return_value=fake_thread,
    )

    with patch.object(runner, "setup_interactive") as mock_setup:
        with threadpool_mock as mock_threadpool:
            runner.run_interactive("fake cmd", "fake host")

    mock_setup.assert_called_once_with("fake host")
    mock_threadpool.assert_called_once_with(max_workers=14)

    stdout, stderr = capfd.readouterr()
    assert stdout == "fake host: fake results\n"
    assert stderr == ""

def test_run_single_success():
    """Ensure the calls to send commands on a single host."""

    runner = Bladerunner({
        "username": "dudeguybro",
        "password": "hunter40",
        "progressbar": True,
        "port": 2012,
    })

    runner.progress = ProgressBar(1, runner.options)

    with patch.object(runner, "connect", return_value=("ok", 0)) as p_connect:
        with patch.object(runner, "send_commands", return_value=[]) as p_send:
            with patch.object(runner, "close") as p_close:
                with patch.object(runner.progress, "update") as p_update:
                    runner._run_single("nowhere")

    p_connect.assert_called_once_with(
        "nowhere",
        "dudeguybro",
        "hunter40",
        2012,
    )
    p_send.assert_called_once_with("ok", "nowhere")
    p_close.assert_called_once_with("ok", True)
    # if the progressbar is used we should tick it once for the check run
    assert p_update.called

def test_connect_from_jb_denied():
    """Ensure we look through the before text for 'Permission denied's."""

    runner = Bladerunner({"jump_host": "mocked"})
    runner.sshc = Mock()
    runner.sshc.before.find = Mock(return_value=1)

    with patch.object(base, "can_resolve", return_value=True):
        with patch.object(runner, "send_interrupt") as p_interrupt:
            ret = runner.connect("home", "self", "hunter22", 443)

    if sys.version_info > (3,):
        expected_call = bytes("Permission denied", "utf-8")
    else:
        expected_call = "Permission denied"

    runner.sshc.before.find.assert_called_once_with(expected_call)
    runner.sshc.sendline.assert_called_once_with("ssh -p 443 -t [email protected]")
    runner.sshc.expect.assert_called_once_with(
        runner.options["passwd_prompts"] +
        runner.options["shell_prompts"] +
        runner.options["extra_prompts"],
        runner.options["timeout"],
    )
    p_interrupt.assert_called_once_with(runner.sshc)
    assert ret == (None, -4)

def test_run_safely_to_parralel():
    """Ensure after the first success the rest are run in parallel."""

    runner = Bladerunner({
        "username": "broguy",
        "password": "hunter14",
        "progressbar": True,
        "port": 2244,
    })

    # we have to set up the progressbar ourselves here, normally it'd happen in
    # run once we know the len of servers to run on.
    runner.progress = ProgressBar(3, runner.options)

    with patch.object(runner, "connect", return_value=("ok", 0)) as p_connect:
        with patch.object(runner, "send_commands", return_value=[]) as p_send:
            with patch.object(runner, "close") as p_close:
                with patch.object(runner, "_run_parallel_no_check") as p_run:
                    with patch.object(runner.progress, "update") as p_update:
                        runner._run_parallel_safely(["1st", "2nd", "3rd"])

    p_connect.assert_called_once_with(
        "1st",
        "broguy",
        "hunter14",
        2244,
    )
    p_send.assert_called_once_with("ok", "1st")
    p_close.assert_called_once_with("ok", True)
    # if the progressbar is used we should update it once for the check run
    assert p_update.called
    p_run.assert_called_once_with(["2nd", "3rd"])

def test_send_cmd_unix_endings(unicode_chr):
    """Ensure the correct line ending is used when unix is specified."""

    runner = Bladerunner({
        "unix_line_endings": True,
        "windows_line_endings": False,
        "second_password": "hunter55",
    })
    server = Mock()
    # server.expect returns an integer of the prompt matched in its list
    # we want to return N+1 to simulate matching a passwd prompt
    server.expect = Mock(return_value=(
        len(runner.options["shell_prompts"]) +
        len(runner.options["extra_prompts"]) +
        1
    ))

    with patch.object(base, "format_output") as p_format_out:
        runner._send_cmd("fake", server)

    p_format_out.assert_called_once_with(server.before, "fake", runner.options)
    server.send.assert_called_once_with("fake{0}".format(unicode_chr(0x000A)))

    # the second password should be send with sendline
    server.sendline.assert_called_once_with("hunter55")

    assert server.expect.call_count == 2

def test_close_and_terminate():
    """Sends 'exit' and terminates the pexpect connection object."""

    runner = Bladerunner()
    sshc = Mock()
    runner.close(sshc, True)
    sshc.sendline.assert_called_once_with("exit")
    assert sshc.terminate.called

def test_login_new_host():
    """The first passwd prompt is a match for a new ssh key identity."""

    runner = Bladerunner()
    sshc = Mock()
    sshc.expect = Mock(side_effect=iter([2, 22]))  # passwd, then shell
    assert runner.login(sshc, "fake", 0) == (sshc, 1)
    assert sshc.sendline.mock_calls == [call("yes"), call("fake")]

def test_login_send_password():
    """Ensure then calls when logging in properly with a passwd prompt."""

    runner = Bladerunner()
    sshc = Mock()
    sshc.expect = Mock(return_value=22)
    assert runner.login(sshc, "mock word", 1) == (sshc, 1)
    sshc.sendline.assert_called_once_with("mock word")

def test_multipass():
    """Ensure the correct calls are made to attempt multiple passwords."""

    runner = Bladerunner()
    sshc = Mock()
    with patch.object(runner, "login", return_value=("fake", 1)) as p_login:
        assert runner._multipass(sshc, "hunter11", 123) == ("fake", 1)

    p_login.assert_called_once_with(sshc, "hunter11", 123)

def test_run_no_options():
    """Uses Mock to enforce the Bladerunner.run logic with no configuration."""

    runner = Bladerunner()

    with patch.object(runner, "_run_parallel") as patched_run:
        runner.run("nothing", "nowhere")

    patched_run.assert_called_once_with(["nowhere"])

def test_serial_execution():
    """If the delay optios is set, we should call _run_serial."""

    runner = Bladerunner({"delay": 10})

    with patch.object(runner, "_run_serial") as patched_run:
        runner.run("nothing", "nowhere")

    patched_run.assert_called_once_with(["nowhere"])

def test_login_unexpected_prompt():
    """We received a password prompt when no password is in use."""

    runner = Bladerunner()
    sshc = Mock()

    with patch.object(runner, "send_interrupt") as p_interrupt:
        assert runner.login(sshc, None, 1) == (None, -2)

    p_interrupt.assert_called_once_with(sshc)