def test_reboot_with_command(m_get_stdout): instance = mock.Mock() instance.private_ip_address = "ip1" instance.tags = {"Name": "n1", "moz-type": "t1"} ssh_client = SSHClient(instance, "u1", "k1") ssh_client.reboot("cmd1") m_get_stdout.assert_called_once_with("cmd1")
def test_connect_returns_None(m_connect): instance = mock.Mock() instance.private_ip_address = "ip1" instance.tags = {"Name": "n1"} ssh_client = SSHClient(instance, "u1", "k1") m_connect.side_effect = Exception("Ooops") assert ssh_client.connect() is None
def test_connect(m_connect): instance = mock.Mock() instance.private_ip_address = "ip1" instance.tags = {"Name": "n1"} ssh_client = SSHClient(instance, "u1", "k1") ssh_client.connect() m_connect.assert_called_once_with(hostname="ip1", username="******", key_filename="k1", timeout=10)
def test_get_stdout(m_exec_command): instance = mock.Mock() instance.private_ip_address = "ip1" instance.tags = {"Name": "n1"} ssh_client = SSHClient(instance, "u1", "k1") stdin, stdout = mock.Mock(), mock.Mock() stdout.read.return_value = "out1" m_exec_command.return_value = stdin, stdout, None out = ssh_client.get_stdout("my command") m_exec_command.assert_called_once_with("my command") stdin.close.assert_called_once_with() stdout.read.assert_called_once_with() assert out == "out1"
def test_policy(): instance = mock.MagicMock() client = SSHClient(instance, "user", "key") assert isinstance(client._policy, paramiko.MissingHostKeyPolicy)
def aws_safe_stop_instance(i, impaired_ids, user, key_filename, masters_json, dryrun=False): "Returns True if stopped" # TODO: Check with slavealloc ssh_client = SSHClient(instance=i, username=user, key_filename=key_filename).connect() stopped = False launch_time = calendar.timegm( time.strptime(i.launch_time[:19], '%Y-%m-%dT%H:%M:%S')) if not ssh_client: if i.id in impaired_ids: if time.time() - launch_time > 60 * 10: stopped = True if not dryrun: log.debug( "%s - shut down an instance with impaired status", ssh_client.name) i.terminate() gr_log.add("impaired.{moz_type}".format( ssh_client.instance.tags.get("moz-type", "none")), 1, collect=True) else: log.debug("%s - would have stopped", ssh_client.name) return stopped uptime_min = int((time.time() - launch_time) / 60) # Don't try to stop spot instances until after STOP_THRESHOLD_MINS_SPOT # minutes into each hour if i.spot_instance_request_id: threshold = STOP_THRESHOLD_MINS_SPOT if uptime_min % 60 < threshold: log.debug("Skipping %s, with uptime %s", ssh_client.name, uptime_min) return False else: # On demand instances can be stopped after STOP_THRESHOLD_MINS_ONDEMAND threshold = STOP_THRESHOLD_MINS_ONDEMAND if uptime_min < threshold: log.debug("Skipping %s, with updtime %s", ssh_client.name, uptime_min) return False last_activity = get_last_activity(ssh_client) if last_activity == ACTIVITY_STOPPED: stopped = True if not dryrun: log.debug("%s - stopping instance (launched %s)", ssh_client.name, i.launch_time) i.terminate() else: log.debug("%s - would have stopped", ssh_client.name) return stopped if last_activity == ACTIVITY_BOOTING: # Wait harder return stopped log.debug("%s - last activity %s", ssh_client.name, last_activity) # If it looks like we're idle for more than 8 hours, kill the machine if last_activity > 8 * 3600: log.debug("%s - last activity more than 8 hours ago; shutting down", ssh_client.name) if not dryrun: log.debug("%s - starting graceful shutdown", ssh_client.name) graceful_shutdown(ssh_client, masters_json) # Stop the instance log.debug("%s - stopping instance", ssh_client.name) i.terminate() stopped = True # If the machine is idle for more than 5 minutes, shut it down elif last_activity > 300: if not dryrun: # Hit graceful shutdown on the master log.debug("%s - starting graceful shutdown", ssh_client.name) graceful_shutdown(ssh_client, masters_json) # Check if we've exited right away if get_last_activity(ssh_client) == ACTIVITY_STOPPED: log.debug("%s - stopping instance", ssh_client.name) i.terminate() stopped = True else: log.debug("%s - not stopping, waiting for graceful shutdown", ssh_client.name) else: log.debug("%s - would have started graceful shutdown", ssh_client.name) stopped = True else: log.debug("%s - not stopping", ssh_client.name) return stopped