def test_simple_successful_kill_job(self):
"""Run a test of the "kill" command against a mocked-out API:
Verifies that the kill command sends the right API RPCs, and performs the correct
tests on the result."""
mock_options = self.setup_mock_options()
mock_config = Mock()
mock_api_factory = self.setup_mock_api_factory()
with contextlib.nested(
patch('twitter.aurora.client.commands.core.make_client_factory',
return_value=mock_api_factory),
patch('twitter.common.app.get_options', return_value=mock_options),
patch('twitter.aurora.client.commands.core.get_job_config', return_value=mock_config)) as (
mock_make_client_factory,
options, mock_get_job_config):
mock_api = mock_api_factory.return_value
with temporary_file() as fp:
fp.write(self.get_valid_config())
fp.flush()
kill(['west/mchucarroll/test/hello', fp.name], mock_options)
# Now check that the right API calls got made.
self.assert_kill_job_called(mock_api)
mock_api.kill_job.assert_called_with(
AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV,
name=self.TEST_JOB), None, config=mock_config)
self.assert_scheduler_called(mock_api)
assert mock_make_client_factory.call_count == 1
def disambiguate_args_or_die(cls,
args,
options,
client_factory=AuroraClientAPI):
"""
Returns a (AuroraClientAPI, AuroraJobKey, AuroraConfigFile:str) tuple
if one can be found given the args, potentially querying the scheduler with the returned client.
Calls die() with an appropriate error message otherwise.
Arguments:
args: args from app command invocation.
options: options from app command invocation. must have env and cluster attributes.
client_factory: a callable (cluster) -> AuroraClientAPI.
"""
if not len(args) > 0:
die('job path is required')
try:
job_key = AuroraJobKey.from_path(args[0])
client = client_factory(job_key.cluster)
config_file = args[1] if len(
args) > 1 else None # the config for hooks
return client, job_key, config_file
except AuroraJobKey.Error:
log.warning(
"Failed to parse job path, falling back to compatibility mode")
role = args[0] if len(args) > 0 else None
name = args[1] if len(args) > 1 else None
env = None
config_file = None # deprecated form does not support hooks functionality
cluster = options.cluster
if not cluster:
die('cluster is required')
client = client_factory(cluster)
return client, cls._disambiguate_or_die(client, role, env,
name), config_file
def setUp(self):
self.RETURN_VALUE = 'foo'
test_obj = self
class FakeAuroraClientAPI(object):
def cancel_update(self, job_key):
test_obj.API_CALL = functools.partial(self.cancel_update, job_key)
return test_obj.RETURN_VALUE
def kill_job(self, job_key, instances=None, lock=None):
test_obj.API_CALL = functools.partial(self.kill_job, job_key, instances, lock)
return test_obj.RETURN_VALUE
def restart(self, job_key, shards, updater_config, health_check_interval_seconds):
test_obj.API_CALL = functools.partial(self.restart, job_key, shards,
updater_config, health_check_interval_seconds)
return test_obj.RETURN_VALUE
def start_cronjob(self, job_key):
test_obj.API_CALL = functools.partial(self.start_cronjob, job_key)
return test_obj.RETURN_VALUE
self._patch_bases(NonHookedAuroraClientAPI, (FakeAuroraClientAPI, ))
self.api = NonHookedAuroraClientAPI()
# Test args passed in to check that these are proxied un-modified
self.test_job_key = AuroraJobKey.from_path('a/b/c/d')
self.test_config = 'bar'
self.test_shards = 'baz'
self.test_lock = 'lock'
self.test_updater_config = 'blah'
self.health_check_interval_seconds = 'baa'
def get_job_config(job_spec, config_file, options):
try:
job_key = AuroraJobKey.from_path(job_spec)
select_cluster = job_key.cluster
select_env = job_key.env
select_role = job_key.role
jobname = job_key.name
except AuroraJobKey.Error:
deprecation_warning('Please refer to your job in CLUSTER/ROLE/ENV/NAME format.')
select_cluster = options.cluster if options.cluster else None
select_env = options.env
select_role = None
jobname = job_spec
try:
json_option = options.json
except AttributeError:
json_option = False
try:
bindings = options.bindings
except AttributeError:
bindings = ()
return get_config(
jobname,
config_file,
json_option,
bindings,
select_cluster=select_cluster,
select_role=select_role,
select_env=select_env)
def test_disambiguate_args_or_die_unambiguous_with_no_config(self):
expected = (self._api,
AuroraJobKey(self.CLUSTER.name, self.ROLE, self.ENV,
self.NAME), None)
result = LiveJobDisambiguator.disambiguate_args_or_die(
[self.JOB_PATH], None, client_factory=lambda *_: self._api)
assert result == expected
def assert_cancel_update_called(cls, mock_api):
# Running cancel update should result in calling the API cancel_update
# method once, with an AuroraJobKey parameter.
assert mock_api.cancel_update.call_count == 1
assert mock_api.cancel_update.called_with(AuroraJobKey(
cls.TEST_CLUSTER, cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB),
config=None)
def run(args, options):
"""usage: run cluster/role/env/job cmd
Runs a shell command on all machines currently hosting shards of a single job.
This feature supports the same command line wildcards that are used to
populate a job's commands.
This means anything in the {{mesos.*}} and {{thermos.*}} namespaces.
"""
# TODO(William Farner): Add support for invoking on individual shards.
# TODO(Kevin Sweeney): Restore the ability to run across jobs with globs (See MESOS-3010).
if not args:
die('job path is required')
job_path = args.pop(0)
try:
cluster_name, role, env, name = AuroraJobKey.from_path(job_path)
except AuroraJobKey.Error as e:
die('Invalid job path "%s": %s' % (job_path, e))
command = ' '.join(args)
cluster = CLUSTERS[cluster_name]
dcr = DistributedCommandRunner(cluster, role, env, [name],
options.ssh_user)
dcr.run(command,
parallelism=options.num_threads,
executor_sandbox=options.executor_sandbox)
def disambiguate_args_or_die(cls, args, options, client_factory=AuroraClientAPI):
"""
Returns a (AuroraClientAPI, AuroraJobKey, AuroraConfigFile:str) tuple
if one can be found given the args, potentially querying the scheduler with the returned client.
Calls die() with an appropriate error message otherwise.
Arguments:
args: args from app command invocation.
options: options from app command invocation. must have env and cluster attributes.
client_factory: a callable (cluster) -> AuroraClientAPI.
"""
if not len(args) > 0:
die('job path is required')
try:
job_key = AuroraJobKey.from_path(args[0])
client = client_factory(job_key.cluster)
config_file = args[1] if len(args) > 1 else None # the config for hooks
return client, job_key, config_file
except AuroraJobKey.Error:
log.warning("Failed to parse job path, falling back to compatibility mode")
role = args[0] if len(args) > 0 else None
name = args[1] if len(args) > 1 else None
env = None
config_file = None # deprecated form does not support hooks functionality
cluster = options.cluster
if not cluster:
die('cluster is required')
client = client_factory(cluster)
return client, cls._disambiguate_or_die(client, role, env, name), config_file
def query_matches(self):
resp = self._client.get_jobs(self._role)
check_and_log_response(resp)
return set(
AuroraJobKey(self._client.cluster.name, j.key.role,
j.key.environment, j.key.name)
for j in resp.result.getJobsResult.configs
if j.key.name == self._name)
def ssh(args, options):
"""usage: ssh cluster/role/env/job shard [args...]
Initiate an SSH session on the machine that a shard is running on.
"""
if not args:
die('Job path is required')
job_path = args.pop(0)
try:
cluster_name, role, env, name = AuroraJobKey.from_path(job_path)
except AuroraJobKey.Error as e:
die('Invalid job path "%s": %s' % (job_path, e))
if not args:
die('Shard is required')
try:
shard = int(args.pop(0))
except ValueError:
die('Shard must be an integer')
api = make_client(cluster_name)
resp = api.query(api.build_query(role, name, set([int(shard)]), env=env))
check_and_log_response(resp)
first_task = resp.result.scheduleStatusResult.tasks[0]
remote_cmd = 'bash' if not args else ' '.join(args)
command = DistributedCommandRunner.substitute(remote_cmd, first_task,
api.cluster, executor_sandbox=options.executor_sandbox)
ssh_command = ['ssh', '-t']
role = first_task.assignedTask.task.owner.role
slave_host = first_task.assignedTask.slaveHost
for tunnel in options.tunnels:
try:
port, name = tunnel.split(':')
port = int(port)
except ValueError:
die('Could not parse tunnel: %s. Must be of form PORT:NAME' % tunnel)
if name not in first_task.assignedTask.assignedPorts:
die('Task %s has no port named %s' % (first_task.assignedTask.taskId, name))
ssh_command += [
'-L', '%d:%s:%d' % (port, slave_host, first_task.assignedTask.assignedPorts[name])]
ssh_command += ['%s@%s' % (options.ssh_user or role, slave_host), command]
return subprocess.call(ssh_command)
def test_kill_job_with_instances(self):
"""Test kill client-side API logic."""
mock_context = FakeAuroraCommandContext()
with contextlib.nested(
patch('twitter.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
patch('twitter.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)):
api = mock_context.get_api('west')
api.kill_job.return_value = self.get_kill_job_response()
with temporary_file() as fp:
fp.write(self.get_valid_config())
fp.flush()
cmd = AuroraCommandLine()
cmd.execute(['job', 'kill', '--config=%s' % fp.name, '--instances=0,2,4-6',
'west/mchucarroll/test/hello'])
# Now check that the right API calls got made.
assert api.kill_job.call_count == 1
api.kill_job.assert_called_with(AuroraJobKey.from_path('west/mchucarroll/test/hello'),
[0, 2, 4, 5, 6])
def setUp(self):
self.RETURN_VALUE = 'foo'
test_obj = self
class FakeAuroraClientAPI(object):
def cancel_update(self, job_key):
test_obj.API_CALL = functools.partial(self.cancel_update,
job_key)
return test_obj.RETURN_VALUE
def kill_job(self, job_key, instances=None, lock=None):
test_obj.API_CALL = functools.partial(self.kill_job, job_key,
instances, lock)
return test_obj.RETURN_VALUE
def restart(self, job_key, shards, updater_config,
health_check_interval_seconds):
test_obj.API_CALL = functools.partial(
self.restart, job_key, shards, updater_config,
health_check_interval_seconds)
return test_obj.RETURN_VALUE
def start_cronjob(self, job_key):
test_obj.API_CALL = functools.partial(self.start_cronjob,
job_key)
return test_obj.RETURN_VALUE
self._patch_bases(NonHookedAuroraClientAPI, (FakeAuroraClientAPI, ))
self.api = NonHookedAuroraClientAPI()
# Test args passed in to check that these are proxied un-modified
self.test_job_key = AuroraJobKey.from_path('a/b/c/d')
self.test_config = 'bar'
self.test_shards = 'baz'
self.test_lock = 'lock'
self.test_updater_config = 'blah'
self.health_check_interval_seconds = 'baa'
def run(args, options):
"""usage: run cluster/role/env/job cmd
Runs a shell command on all machines currently hosting shards of a single job.
This feature supports the same command line wildcards that are used to
populate a job's commands.
This means anything in the {{mesos.*}} and {{thermos.*}} namespaces.
"""
# TODO(William Farner): Add support for invoking on individual shards.
# TODO(Kevin Sweeney): Restore the ability to run across jobs with globs (See MESOS-3010).
if not args:
die('job path is required')
job_path = args.pop(0)
try:
cluster_name, role, env, name = AuroraJobKey.from_path(job_path)
except AuroraJobKey.Error as e:
die('Invalid job path "%s": %s' % (job_path, e))
command = ' '.join(args)
cluster = CLUSTERS[cluster_name]
dcr = DistributedCommandRunner(cluster, role, env, [name], options.ssh_user)
dcr.run(command, parallelism=options.num_threads, executor_sandbox=options.executor_sandbox)
def _disambiguate_or_die(cls, client, role, env, name):
# Returns a single AuroraJobKey if one can be found given the args, potentially
# querying the scheduler. Calls die() with an appropriate error message otherwise.
try:
disambiguator = cls(client, role, env, name)
except ValueError as e:
die(e)
if not disambiguator.ambiguous:
return AuroraJobKey(client.cluster.name, role, env, name)
deprecation_warning(
"Job ambiguously specified - querying the scheduler to disambiguate"
)
matches = disambiguator.query_matches()
if len(matches) == 1:
(match, ) = matches
log.info("Found job %s" % match)
return match
elif len(matches) == 0:
die("No jobs found")
else:
die("Multiple jobs match (%s) - disambiguate by using the CLUSTER/ROLE/ENV/NAME form"
% ",".join(str(m) for m in matches))
def parse_aurora_job_key_into(option, opt, value, parser):
try:
setattr(parser.values, option.dest, AuroraJobKey.from_path(value))
except AuroraJobKey.Error as e:
raise optparse.OptionValueError('Failed to parse: %s' % e)
def test_basic(self):
AuroraJobKey.from_path("smf1/mesos/test/labrat")
def job_key(self):
return AuroraJobKey(self.cluster(), self.role(), self.environment(),
self.name())
def test_basic(self):
AuroraJobKey.from_path("smf1/mesos/test/labrat")
from twitter.aurora.client.api.restarter import Restarter
from twitter.aurora.client.api.instance_watcher import InstanceWatcher
from twitter.aurora.client.api.updater_util import UpdaterConfig
from twitter.aurora.common.aurora_job_key import AuroraJobKey
from gen.twitter.aurora.AuroraSchedulerManager import Client as scheduler_client
from gen.twitter.aurora.ttypes import *
from mox import IgnoreArg, MoxTestBase
# test space
from twitter.aurora.client.fake_scheduler_proxy import FakeSchedulerProxy
SESSION_KEY = 'test_session'
CLUSTER = 'smfd'
JOB = AuroraJobKey(CLUSTER, 'johndoe', 'test', 'test_job')
HEALTH_CHECK_INTERVAL_SECONDS = 5
UPDATER_CONFIG = UpdaterConfig(
2, # batch_size
23, # restart_threshold
45, #watch_secs
0, # max_per_instance_failures
0 # max_total_failures
)
class TestRestarter(MoxTestBase):
def setUp(self):
super(TestRestarter, self).setUp()
self.mock_scheduler = self.mox.CreateMock(scheduler_client)
