def parse_task_instance_key(key):
pieces = key.split('/')
if len(pieces) != 5:
raise ValueError('Task instance specifier %s is not in the form '
'CLUSTER/ROLE/ENV/NAME/INSTANCE' % key)
(cluster, role, env, name, instance_str) = pieces
try:
instance = int(instance_str)
except ValueError:
raise ValueError('Instance must be an integer, but got %s' % instance_str)
return TaskInstanceKey(AuroraJobKey(cluster, role, env, name), instance)
def test_successful_status_shallow_nometadata(self):
"""Regression test: there was a crasher bug when metadata was None."""
mock_context = FakeAuroraCommandContext()
mock_api = mock_context.get_api('west')
mock_api.check_status.return_value = self.create_status_null_metadata()
with contextlib.nested(
patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
cmd = AuroraCommandLine()
cmd.execute(['job', 'status', 'west/bozo/test/hello'])
mock_api.check_status.assert_called_with(AuroraJobKey('west', 'bozo', 'test', 'hello'))
def test_add_instances(self):
"""Test adding instances."""
api, mock_proxy = self.mock_api()
job_key = AuroraJobKey("foo", "role", "env", "name")
mock_proxy.addInstances.return_value = self.create_simple_success_response(
)
api.add_instances(job_key, 1, 10)
mock_proxy.addInstances.assert_called_once_with(
None, None, InstanceKey(jobKey=job_key.to_thrift(), instanceId=1),
10)
def test_query_job_updates(self):
"""Test querying job updates."""
api, mock_proxy = self.mock_api()
job_key = AuroraJobKey("foo", "role", "env", "name")
query = JobUpdateQuery(
jobKey=job_key.to_thrift(),
updateStatuses={JobUpdateStatus.ROLLING_FORWARD})
api.query_job_updates(job_key=job_key,
update_statuses=query.updateStatuses)
mock_proxy.getJobUpdateSummaries.assert_called_once_with(query,
retry=True)
def test_get_domain_uptime_vector_with_hosts(self):
with patch('apache.aurora.client.api.sla.task_query', return_value=TaskQuery()) as (mock_query):
self.mock_get_tasks([
self.create_task(100, 1, 'h1', 'j1'),
self.create_task(200, 2, 'h1', 'j2'),
self.create_task(200, 3, 'h2', 'j1'),
self.create_task(200, 3, 'h2', 'j3'),
self.create_task(200, 4, 'h3', 'j4'),
self.create_task(200, 4, 'h3', 'j3'),
])
hosts = ['h1', 'h2', 'h3']
jobs = set([
AuroraJobKey(self._cluster.name, self._role, self._env, 'j1'),
AuroraJobKey(self._cluster.name, self._role, self._env, 'j2'),
AuroraJobKey(self._cluster.name, self._role, self._env, 'j3'),
AuroraJobKey(self._cluster.name, self._role, self._env, 'j4')
])
self._sla.get_domain_uptime_vector(self._cluster, self._min_count, hosts)
mock_query.assert_has_calls([call(hosts=hosts), call(job_keys=jobs)], any_order=False)
def job_key_from_scheduled(task, cluster):
"""Creates AuroraJobKey from the ScheduledTask.
Arguments:
task -- ScheduledTask to get job key from.
cluster -- Cluster the task belongs to.
"""
return AuroraJobKey(cluster=cluster.name,
role=task.assignedTask.task.owner.role,
env=task.assignedTask.task.environment,
name=task.assignedTask.task.jobName)
def test_successful_status_shallow(self):
"""Test the status command at the shallowest level: calling status should end up invoking
the local APIs get_status method."""
mock_context = FakeAuroraCommandContext()
mock_api = mock_context.get_api('west')
mock_api.check_status.return_value = self.create_status_response()
with contextlib.nested(
patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
cmd = AuroraCommandLine()
cmd.execute(['job', 'status', 'west/bozo/test/hello'])
mock_api.check_status.assert_called_with(AuroraJobKey('west', 'bozo', 'test', 'hello'))
def get_job_list(self, clusters, role=None):
"""Get a list of jobs from a group of clusters.
:param clusters: the clusters to query for jobs
:param role: if specified, only return jobs for the role; otherwise, return all jobs.
"""
result = []
if '*' in role:
role = None
for cluster in clusters:
api = self.get_api(cluster)
resp = api.get_jobs(role)
self.log_response_and_raise(resp, err_code=EXIT_COMMAND_FAILURE)
result.extend([AuroraJobKey(cluster, job.key.role, job.key.environment, job.key.name)
for job in resp.result.getJobsResult.configs])
return result
def test_happy_hook(self):
"""Test that hooks that return 0 don't block command execution"""
class HappyHook(CoreCommandHook):
@property
def name(self):
return "I'm so happy"
def execute(self, cmd, options, *args, **kwargs):
return 0
CoreCommandHook.register_hook(HappyHook())
mock_options = self.setup_mock_options()
mock_config = Mock()
(mock_api, mock_scheduler_proxy) = self.create_mock_api()
mock_api.kill_job.return_value = self.get_kill_job_response()
mock_scheduler_proxy.killTasks.return_value = self.get_kill_job_response()
mock_query_results = [
self.create_mock_status_query_result(ScheduleStatus.RUNNING),
self.create_mock_status_query_result(ScheduleStatus.KILLING),
self.create_mock_status_query_result(ScheduleStatus.KILLED),
]
mock_scheduler_proxy.getTasksWithoutConfigs.side_effect = mock_query_results
with contextlib.nested(
patch('threading._Event.wait'),
patch('apache.aurora.client.commands.core.make_client',
return_value=mock_api),
patch('twitter.common.app.get_options', return_value=mock_options),
patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)) as (
sleep,
mock_make_client,
options,
mock_get_job_config):
with temporary_file() as fp:
fp.write(self.get_valid_config())
fp.flush()
killall(['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, self.get_expected_task_query(), 3)
CoreCommandHook.clear_hooks()
def test_restart_inactive_instance_spec(self):
command = RestartCommand()
jobkey = AuroraJobKey("cluster", "role", "env", "job")
mock_options = mock_verb_options(command)
mock_options.instance_spec = TaskInstanceKey(jobkey, [1])
mock_options.strict = True
fake_context = FakeAuroraCommandContext()
fake_context.set_options(mock_options)
fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())
with pytest.raises(Context.CommandError) as e:
command.execute(fake_context)
assert e.message == "Invalid instance parameter: [1]"
def test_status_wildcard_two(self):
"""Test status using a wildcard. It should first call api.get_jobs, and then do a
getTasksWithoutConfigs on each job. This time, use a pattern that doesn't match
all of the jobs."""
mock_context = FakeAuroraCommandContext()
mock_api = mock_context.get_api('west')
mock_api.check_status.return_value = self.create_status_response()
mock_api.get_jobs.return_value = self.create_getjobs_response()
with contextlib.nested(
patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
cmd = AuroraCommandLine()
cmd.execute(['job', 'status', 'example/*/*/hello'])
# Wildcard should have expanded to two jobs, but only matched one,
# so there should be one call to check_status.
assert mock_api.check_status.call_count == 1
mock_api.check_status.assert_called_with(
AuroraJobKey('example', 'bozo', 'test', 'hello'))
def test_kill_batched_queries_active_instances(self):
"""Verify that the batch kill operates on active instances only."""
command = KillCommand()
jobkey = AuroraJobKey("cluster", "role", "env", "job")
mock_options = mock_verb_options(command)
mock_options.instance_spec = TaskInstanceKey(jobkey, [1])
mock_options.no_batching = False
fake_context = FakeAuroraCommandContext()
fake_context.set_options(mock_options)
fake_context.add_expected_query_result(
AuroraClientCommandTest.create_empty_task_result())
command.execute(fake_context)
assert fake_context.get_err(
)[0] == "No tasks to kill found for job cluster/role/env/job"
def test_handles_api_auth_error():
context = AuroraCommandContext()
mock_scheduler_proxy = mock.create_autospec(spec=SchedulerProxyApiSpec,
instance=True)
mock_scheduler_proxy.killTasks.side_effect = SchedulerProxy.AuthError()
mock_api = AuroraClientAPI(TEST_CLUSTER, 'user-agent')
mock_api._scheduler_proxy = mock_scheduler_proxy
context.apis = {TEST_CLUSTER.name: mock_api}
api = context.get_api(TEST_CLUSTER.name,
clusters={TEST_CLUSTER.name: TEST_CLUSTER})
with pytest.raises(Context.CommandError) as e:
api.kill_job(AuroraJobKey(TEST_CLUSTER.name, 'role', 'env', 'job'))
assert e.value.code == EXIT_AUTH_ERROR
assert mock_scheduler_proxy.killTasks.call_count == 1
def test_restart_opens_url(self):
command = RestartCommand()
jobkey = AuroraJobKey("cluster", "role", "env", "job")
mock_options = mock_verb_options(command)
mock_options.instance_spec = TaskInstanceKey(jobkey, None)
mock_options.strict = True
mock_options.open_browser = True
fake_context = FakeAuroraCommandContext()
fake_context.set_options(mock_options)
fake_api = fake_context.fake_api
fake_api.restart.return_value = AuroraClientCommandTest.create_simple_success_response()
command.execute(fake_context)
assert self.mock_webbrowser.mock_calls == [
call("http://something_or_other/scheduler/role/env/job")
]
def test_successful_schedule(self):
mock_context = FakeAuroraCommandContext()
key = AuroraJobKey("west", "bozo", "test", "hello")
with contextlib.nested(
patch('apache.aurora.client.cli.cron.CronNoun.create_context', return_value=mock_context)):
api = mock_context.get_api('west')
api.schedule_cron.return_value = self.create_simple_success_response()
with temporary_file() as fp:
fp.write(self.get_valid_cron_config())
fp.flush()
cmd = AuroraCommandLine()
cmd.execute(['cron', 'schedule', key.to_path(), fp.name])
# Now check that the right API calls got made.
# Check that create_job was called exactly once, with an AuroraConfig parameter.
assert api.schedule_cron.call_count == 1
assert isinstance(api.schedule_cron.call_args[0][0], AuroraConfig)
# The last text printed out to the user should contain a url to the job
assert mock_context.get_job_page(api, key) in mock_context.out[-1]
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 get_update_details_response(cls):
query_response = Response()
query_response.responseCode = ResponseCode.OK
query_response.result = Result()
details = JobUpdateDetails()
query_response.result.getJobUpdateDetailsResult = GetJobUpdateDetailsResult(
details=details)
details.update = JobUpdate()
details.update.summary = JobUpdateSummary(
jobKey=AuroraJobKey('west', 'mcc', 'test', 'hello'),
updateId="fake-update-identifier",
user="******",
state=JobUpdateState(status=JobUpdateStatus.ROLLING_FORWARD,
createdTimestampMs=1411404927,
lastModifiedTimestampMs=14114056030))
details.updateEvents = [
JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD,
timestampMs=1411404927),
JobUpdateEvent(status=JobUpdateStatus.ROLL_FORWARD_PAUSED,
timestampMs=1411405000),
JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD,
timestampMs=1411405100)
]
details.instanceEvents = [
JobInstanceUpdateEvent(instanceId=1,
timestampMs=1411404930,
action=JobUpdateAction.INSTANCE_UPDATING),
JobInstanceUpdateEvent(instanceId=2,
timestampMs=1411404940,
action=JobUpdateAction.INSTANCE_UPDATING),
JobInstanceUpdateEvent(instanceId=1,
timestampMs=1411404950,
action=JobUpdateAction.INSTANCE_UPDATED),
JobInstanceUpdateEvent(instanceId=2,
timestampMs=1411404960,
action=JobUpdateAction.INSTANCE_UPDATED)
]
return query_response
def test_kill_lock_error_nobatch(self):
"""Verify that the no batch code path correctly includes the lock error message."""
command = KillCommand()
jobkey = AuroraJobKey("cluster", "role", "env", "job")
mock_options = mock_verb_options(command)
mock_options.instance_spec = TaskInstanceKey(jobkey, [])
mock_options.no_batching = True
fake_context = FakeAuroraCommandContext()
fake_context.set_options(mock_options)
mock_api = fake_context.get_api('test')
mock_api.kill_job.return_value = AuroraClientCommandTest.create_blank_response(
ResponseCode.LOCK_ERROR, "Error.")
with pytest.raises(Context.CommandError):
command.execute(fake_context)
mock_api.kill_job.assert_called_once_with(
jobkey, mock_options.instance_spec.instance)
self.assert_lock_message(fake_context)
from apache.aurora.client.api.instance_watcher import InstanceWatcher
from apache.aurora.client.api.restarter import Restarter
from apache.aurora.client.api.updater_util import UpdaterConfig
from apache.aurora.client.fake_scheduler_proxy import FakeSchedulerProxy
from apache.aurora.common.aurora_job_key import AuroraJobKey
from gen.apache.aurora.api.AuroraSchedulerManager import Client as scheduler_client
from gen.apache.aurora.api.ttypes import (AssignedTask, Response, ResponseCode,
Result, ScheduledTask,
ScheduleStatus, ScheduleStatusResult,
TaskConfig)
SESSION_KEY = 'test_session'
CLUSTER = 'smfd'
JOB = AuroraJobKey(CLUSTER, 'johndoe', 'test', 'test_job')
HEALTH_CHECK_INTERVAL_SECONDS = 5
UPDATER_CONFIG = UpdaterConfig(
batch_size=2,
restart_threshold=23,
watch_secs=45,
max_per_shard_failures=0,
max_total_failures=0,
rollback_on_failure=True,
)
class TestRestarter(MoxTestBase):
def setUp(self):
super(TestRestarter, self).setUp()
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)
class AuroraClientCommandTest(unittest.TestCase):
FAKE_TIME = 42131
@classmethod
def create_blank_response(cls, code, msg):
return Response(
responseCode=code,
details=[ResponseDetail(message=msg)],
serverInfo=ServerInfo(thriftAPIVersion=CURRENT_API_VERSION.major))
@classmethod
def create_simple_success_response(cls):
return cls.create_blank_response(ResponseCode.OK, 'OK')
@classmethod
def create_error_response(cls):
return cls.create_blank_response(ResponseCode.ERROR, 'Whoops')
@classmethod
def create_mock_api(cls):
"""Builds up a mock API object, with a mock SchedulerProxy"""
mock_scheduler = create_autospec(spec=SchedulerThriftApiSpec,
instance=True)
mock_scheduler.url = "http://something_or_other"
mock_scheduler_client = create_autospec(spec=SchedulerProxyApiSpec,
instance=True)
mock_scheduler_client.url = "http://something_or_other"
mock_api = create_autospec(spec=HookedAuroraClientAPI, instance=True)
mock_api.scheduler_proxy = mock_scheduler_client
return mock_api, mock_scheduler_client
@classmethod
def create_mock_api_factory(cls):
"""Create a collection of mocks for a test that wants to mock out the client API
by patching the api factory."""
mock_api, mock_scheduler_client = cls.create_mock_api()
mock_api_factory = lambda: mock_api
return mock_api_factory, mock_scheduler_client
@classmethod
def create_query_call_result(cls, task=None):
status_response = cls.create_empty_task_result()
if task is None:
for i in range(20):
status_response.result.scheduleStatusResult.tasks.append(
cls.create_scheduled_task(i))
else:
status_response.result.scheduleStatusResult.tasks.append(task)
return status_response
@classmethod
def create_empty_task_result(cls):
status_response = cls.create_simple_success_response()
status_response.result = Result(
scheduleStatusResult=ScheduleStatusResult(tasks=[]))
return status_response
@classmethod
def create_scheduled_task(cls,
instance_id,
status=ScheduleStatus.RUNNING,
task_id=None,
initial_time=None):
task = ScheduledTask(
status=status,
assignedTask=AssignedTask(instanceId=instance_id,
taskId=task_id or "Task%s" % instance_id,
slaveId="Slave%s" % instance_id,
slaveHost="Slave%s" % instance_id,
task=TaskConfig()),
taskEvents=[TaskEvent(timestamp=initial_time or 1000)])
return task
@classmethod
def create_scheduled_tasks(cls):
tasks = []
for name in ['foo', 'bar', 'baz']:
task = ScheduledTask(
failureCount=0,
assignedTask=AssignedTask(
taskId=1287391823,
slaveHost='slavehost',
task=TaskConfig(
maxTaskFailures=1,
executorConfig=ExecutorConfig(data='fake data'),
metadata=[],
job=JobKey(role=cls.TEST_ROLE,
environment=cls.TEST_ENV,
name=name),
owner=Identity(role=cls.TEST_ROLE),
environment=cls.TEST_ENV,
jobName=name,
numCpus=2,
ramMb=2,
diskMb=2),
instanceId=4237894,
assignedPorts={}),
status=ScheduleStatus.RUNNING,
taskEvents=[
TaskEvent(timestamp=28234726395,
status=ScheduleStatus.RUNNING,
message="Hi there")
])
tasks.append(task)
return tasks
@classmethod
def setup_get_tasks_status_calls(cls, scheduler):
status_response = cls.create_query_call_result()
scheduler.getTasksWithoutConfigs.return_value = status_response
@classmethod
def fake_time(cls, ignored):
"""Utility function used for faking time to speed up tests."""
cls.FAKE_TIME += 2
return cls.FAKE_TIME
CONFIG_BASE = """
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '%(cluster)s',
environment = '%(env)s',
instances = 20,
%(inner)s
update_config = UpdateConfig(
batch_size = 1,
restart_threshold = 60,
watch_secs = 45,
max_per_shard_failures = 2,
),
task = Task(
name = 'test',
processes = [Process(name = 'hello_world', cmdline = 'echo {{thermos.ports[http]}}')],
resources = Resources(cpu = 0.1, ram = 64 * MB, disk = 64 * MB),
)
)
jobs = [HELLO_WORLD]
"""
CRON_CONFIG_BASE = """
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '%(cluster)s',
environment = '%(env)s',
cron_schedule = '*/5 * * * *',
%(inner)s
task = SimpleTask('test', 'echo test')
)
jobs = [HELLO_WORLD]
"""
UNBOUND_CONFIG = textwrap.dedent("""\
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '{{cluster_binding}}',
environment = '%(env)s',
instances = '{{instances_binding}}',
update_config = UpdateConfig(
batch_size = "{{TEST_BATCH}}",
restart_threshold = 60,
watch_secs = 45,
max_per_shard_failures = 2,
),
task = Task(
name = 'test',
processes = [Process(name = 'hello_world', cmdline = 'echo {{thermos.ports[http]}}')],
resources = Resources(cpu = 0.1, ram = 64 * MB, disk = 64 * MB),
)
)
jobs = [HELLO_WORLD]
""")
TEST_ROLE = 'bozo'
TEST_ENV = 'test'
TEST_JOB = 'hello'
TEST_CLUSTER = TEST_CLUSTER
TEST_JOBSPEC = 'west/bozo/test/hello'
TEST_JOBKEY = AuroraJobKey('west', 'bozo', 'test', 'hello')
TEST_CLUSTERS = TEST_CLUSTERS
@classmethod
def get_instance_spec(cls, instances_spec):
"""Create a job instance spec string"""
return '%s/%s' % (cls.TEST_JOBSPEC, instances_spec)
@classmethod
def get_test_config(cls, base, cluster, role, env, job, inner=''):
"""Create a config from the template"""
return base % {
'job': job,
'role': role,
'env': env,
'cluster': cluster,
'inner': inner
}
@classmethod
def get_unbound_test_config(cls, role=None, env=None, job=None):
result = cls.UNBOUND_CONFIG % {
'job': job or cls.TEST_JOB,
'role': role or cls.TEST_ROLE,
'env': env or cls.TEST_ENV
}
return result
@classmethod
def get_valid_config(cls):
return cls.get_test_config(cls.CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB)
@classmethod
def get_valid_cron_config(cls):
return cls.get_test_config(cls.CRON_CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB)
@classmethod
def get_invalid_config(cls, bad_clause):
return cls.get_test_config(cls.CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB,
bad_clause)
@classmethod
def get_invalid_cron_config(cls, bad_clause):
return cls.get_test_config(cls.CRON_CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB,
bad_clause)
@classmethod
def assert_lock_message(cls, context):
assert [
line for line in context.get_err()
if line == "\t%s" % context.LOCK_ERROR_MSG
]
class AuroraClientCommandTest(unittest.TestCase):
FAKE_TIME = 42131
def setUp(self):
patcher = patch('webbrowser.open_new_tab')
self.mock_webbrowser = patcher.start()
self.addCleanup(patcher.stop)
def run(self, result=None):
# Since CLUSTERS is a global value that evaluates code on import this is the best way to
# ensure it does not pollute any tests.
with CLUSTERS.patch(self.TEST_CLUSTERS._clusters.values()):
super(AuroraClientCommandTest, self).run(result)
@classmethod
def create_blank_response(cls, code, msg):
return Response(responseCode=code,
details=[ResponseDetail(message=msg)])
@classmethod
def create_simple_success_response(cls):
return cls.create_blank_response(ResponseCode.OK, 'OK')
@classmethod
def create_error_response(cls):
return cls.create_blank_response(ResponseCode.ERROR, 'Whoops')
@classmethod
def create_mock_api(cls):
"""Builds up a mock API object, with a mock SchedulerProxy"""
mock_scheduler = create_autospec(spec=SchedulerThriftApiSpec,
instance=True)
mock_scheduler.url = "http://something_or_other"
mock_scheduler_client = create_autospec(spec=SchedulerProxyApiSpec,
instance=True)
mock_scheduler_client.url = "http://something_or_other"
mock_api = create_autospec(spec=HookedAuroraClientAPI, instance=True)
mock_api.scheduler_proxy = mock_scheduler_client
return mock_api, mock_scheduler_client
@classmethod
def create_mock_api_factory(cls):
"""Create a collection of mocks for a test that wants to mock out the client API
by patching the api factory."""
mock_api, mock_scheduler_client = cls.create_mock_api()
mock_api_factory = lambda: mock_api
return mock_api_factory, mock_scheduler_client
@classmethod
def create_query_call_result(cls, task=None):
status_response = cls.create_empty_task_result()
if task is None:
for i in range(20):
status_response.result.scheduleStatusResult.tasks.append(
cls.create_scheduled_task(i))
else:
status_response.result.scheduleStatusResult.tasks.append(task)
return status_response
@classmethod
def create_start_job_update_result(cls, code, msg, key, metadata):
resp = cls.create_blank_response(code, msg)
resp.result = Result(startJobUpdateResult=StartJobUpdateResult(
key=key, updateSummary=JobUpdateSummary(metadata=metadata)))
return resp
@classmethod
def create_empty_task_result(cls):
status_response = cls.create_simple_success_response()
status_response.result = Result(
scheduleStatusResult=ScheduleStatusResult(tasks=[]))
return status_response
@classmethod
def create_scheduled_task(cls,
instance_id,
status=ScheduleStatus.RUNNING,
task_id=None,
initial_time=None):
task = ScheduledTask(
status=status,
assignedTask=AssignedTask(instanceId=instance_id,
taskId=task_id or "Task%s" % instance_id,
slaveId="Slave%s" % instance_id,
slaveHost="Slave%s" % instance_id,
task=TaskConfig()),
taskEvents=[TaskEvent(timestamp=initial_time or 1000)])
return task
@classmethod
def create_task_config(cls, name):
return TaskConfig(
maxTaskFailures=1,
executorConfig=ExecutorConfig(data='{"fake": "data"}'),
metadata=[],
job=JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV,
name=name),
numCpus=2,
ramMb=2,
diskMb=2)
@classmethod
def create_scheduled_tasks(cls):
tasks = []
for name in ['foo', 'bar', 'baz']:
task = ScheduledTask(failureCount=0,
assignedTask=AssignedTask(
taskId=1287391823,
slaveHost='slavehost',
task=cls.create_task_config(name),
instanceId=4237894,
assignedPorts={}),
status=ScheduleStatus.RUNNING,
taskEvents=[
TaskEvent(timestamp=28234726395,
status=ScheduleStatus.RUNNING,
message="Hi there")
])
tasks.append(task)
return tasks
@classmethod
def setup_get_tasks_status_calls(cls, scheduler):
status_response = cls.create_query_call_result()
scheduler.getTasksWithoutConfigs.return_value = status_response
@classmethod
def fake_time(cls, ignored):
"""Utility function used for faking time to speed up tests."""
cls.FAKE_TIME += 2
return cls.FAKE_TIME
CONFIG_BASE = """
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '%(cluster)s',
environment = '%(env)s',
instances = 20,
%(inner)s
update_config = UpdateConfig(
batch_size = 1,
watch_secs = 45,
max_per_shard_failures = 2,
),
task = Task(
name = 'test',
processes = [Process(name = 'hello_world', cmdline = 'echo {{thermos.ports[http]}}')],
resources = Resources(cpu = 0.1, ram = 64 * MB, disk = 64 * MB),
)
)
jobs = [HELLO_WORLD]
"""
CRON_CONFIG_BASE = """
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '%(cluster)s',
environment = '%(env)s',
cron_schedule = '*/5 * * * *',
%(inner)s
task = SimpleTask('test', 'echo test')
)
jobs = [HELLO_WORLD]
"""
UNBOUND_CONFIG = textwrap.dedent("""\
HELLO_WORLD = Job(
name = '%(job)s',
role = '%(role)s',
cluster = '{{cluster_binding}}',
environment = '%(env)s',
instances = '{{instances_binding}}',
update_config = UpdateConfig(
batch_size = "{{TEST_BATCH}}",
watch_secs = 45,
max_per_shard_failures = 2,
),
task = Task(
name = 'test',
processes = [Process(
name = 'hello_world',
cmdline = 'echo {{thermos.ports[http]}} {{flags_binding}}'
)],
resources = Resources(cpu = 0.1, ram = 64 * MB, disk = 64 * MB),
)
)
jobs = [HELLO_WORLD]
""")
TEST_ROLE = 'bozo'
TEST_ENV = 'test'
TEST_JOB = 'hello'
TEST_CLUSTER = 'west'
TEST_JOBSPEC = 'west/bozo/test/hello'
TEST_JOBKEY = AuroraJobKey('west', 'bozo', 'test', 'hello')
TEST_CLUSTERS = Clusters([
Cluster(name=TEST_CLUSTER,
zk='zookeeper.example.com',
scheduler_zk_path='/foo/bar',
auth_mechanism='UNAUTHENTICATED')
])
@classmethod
def get_instance_spec(cls, instances_spec):
"""Create a job instance spec string"""
return '%s/%s' % (cls.TEST_JOBSPEC, instances_spec)
@classmethod
def get_test_config(cls, base, cluster, role, env, job, inner=''):
"""Create a config from the template"""
return base % {
'job': job,
'role': role,
'env': env,
'cluster': cluster,
'inner': inner
}
@classmethod
def get_unbound_test_config(cls, role=None, env=None, job=None):
result = cls.UNBOUND_CONFIG % {
'job': job or cls.TEST_JOB,
'role': role or cls.TEST_ROLE,
'env': env or cls.TEST_ENV
}
return result
@classmethod
def get_valid_config(cls):
return cls.get_test_config(cls.CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB)
@classmethod
def get_valid_cron_config(cls):
return cls.get_test_config(cls.CRON_CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB)
@classmethod
def get_invalid_config(cls, bad_clause):
return cls.get_test_config(cls.CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB,
bad_clause)
@classmethod
def get_invalid_cron_config(cls, bad_clause):
return cls.get_test_config(cls.CRON_CONFIG_BASE, cls.TEST_CLUSTER,
cls.TEST_ROLE, cls.TEST_ENV, cls.TEST_JOB,
bad_clause)
@classmethod
def assert_lock_message(cls, context):
assert [
line for line in context.get_err()
if line == "\t%s" % context.LOCK_ERROR_MSG
]
PREFERRED_TIER = TierConfig(name='preferred',
settings={
'preemptible': 'false',
'revocable': 'false'
})
PREEMPTIBLE_TIER = TierConfig(name='preemptible',
settings={
'preemptible': 'true',
'revocable': 'false'
})
REVOCABLE_TIER = TierConfig(name='revocable',
settings={
'preemptible': 'true',
'revocable': 'true'
})
@classmethod
def get_mock_tier_configurations(cls):
response = cls.create_simple_success_response()
response.result = Result(getTierConfigResult=GetTierConfigResult(
defaultTierName=cls.PREEMPTIBLE_TIER.name,
tiers=frozenset([
cls.PREFERRED_TIER, cls.PREEMPTIBLE_TIER, cls.REVOCABLE_TIER
])))
return response
def test_equality(self):
key1 = AuroraJobKey('cluster', 'role', 'env', 'name')
key2 = AuroraJobKey('cluster', 'role', 'env', 'name')
assert key1 == key2
assert not (key1 != key2)
def setUp(self):
self._scheduler = Mock()
self._job_key = AuroraJobKey('cl', 'johndoe', 'test', 'test_job')
self._event = FakeEvent()
class TestJobUpdateApis(unittest.TestCase):
"""Job update APIs tests."""
UPDATE_CONFIG = {
'batch_size': 1,
'watch_secs': 50,
'max_per_shard_failures': 2,
'max_total_failures': 1,
'rollback_on_failure': True,
'wait_for_batch_completion': False,
}
JOB_KEY = AuroraJobKey("foo", "role", "env", "name")
UPDATE_KEY = JobUpdateKey(job=JOB_KEY.to_thrift(), id="id")
@classmethod
def create_blank_response(cls, code, msg):
return Response(responseCode=code,
details=[ResponseDetail(message=msg)],
result=create_autospec(spec=Result,
spec_set=True,
instance=True))
@classmethod
def create_simple_success_response(cls):
return cls.create_blank_response(ResponseCode.OK, 'OK')
@classmethod
def create_error_response(cls):
return cls.create_blank_response(ResponseCode.ERROR, 'ERROR')
@classmethod
def mock_api(cls):
api = AuroraClientAPI(Cluster(name="foo"), 'test-client')
mock_proxy = create_autospec(spec=SchedulerProxyApiSpec,
spec_set=True,
instance=True)
api._scheduler_proxy = mock_proxy
return api, mock_proxy
@classmethod
def create_update_settings(cls):
return JobUpdateSettings(
updateGroupSize=1,
updateStrategy=JobUpdateStrategy(
queueStrategy=QueueJobUpdateStrategy(groupSize=1),
batchStrategy=None,
varBatchStrategy=None),
maxPerInstanceFailures=2,
maxFailedInstances=1,
minWaitInInstanceRunningMs=50 * 1000,
rollbackOnFailure=True,
waitForBatchCompletion=False,
slaAware=False)
@classmethod
def create_update_request(cls, task_config):
return JobUpdateRequest(instanceCount=5,
settings=cls.create_update_settings(),
taskConfig=task_config)
@classmethod
def mock_job_config(cls, error=None):
config = create_autospec(spec=AuroraConfig, instance=True)
update_config = UpdateConfig(batch_size=1,
watch_secs=50,
max_per_shard_failures=2,
max_total_failures=1)
if error:
config.update_config.side_effect = error
else:
config.update_config.return_value = update_config
mock_task_config = create_autospec(spec=JobConfiguration,
instance=True)
mock_task_config.taskConfig = TaskConfig()
config.job.return_value = mock_task_config
config.instances.return_value = 5
return config
def test_add_instances(self):
"""Test adding instances."""
api, mock_proxy = self.mock_api()
job_key = AuroraJobKey("foo", "role", "env", "name")
mock_proxy.addInstances.return_value = self.create_simple_success_response(
)
api.add_instances(job_key, 1, 10)
mock_proxy.addInstances.assert_called_once_with(
InstanceKey(jobKey=job_key.to_thrift(), instanceId=1), 10)
def test_start_job_update(self):
"""Test successful job update start."""
api, mock_proxy = self.mock_api()
task_config = TaskConfig()
mock_proxy.startJobUpdate.return_value = self.create_simple_success_response(
)
api.start_job_update(self.mock_job_config(),
instances=None,
message='hello')
mock_proxy.startJobUpdate.assert_called_once_with(
self.create_update_request(task_config), 'hello', retry=True)
def test_start_job_update_fails_parse_update_config(self):
"""Test start_job_update fails to parse invalid UpdateConfig."""
api, mock_proxy = self.mock_api()
self.assertRaises(AuroraClientAPI.UpdateConfigError,
api.start_job_update,
self.mock_job_config(error=ValueError()), None)
def test_get_job_update_diff(self):
"""Test getting job update diff."""
api, mock_proxy = self.mock_api()
task_config = TaskConfig()
mock_proxy.getJobUpdateDiff.return_value = self.create_simple_success_response(
)
api.get_job_update_diff(self.mock_job_config(), instances=None)
mock_proxy.getJobUpdateDiff.assert_called_once_with(
self.create_update_request(task_config), retry=True)
def test_pause_job_update(self):
"""Test successful job update pause."""
api, mock_proxy = self.mock_api()
mock_proxy.pauseJobUpdate.return_value = self.create_simple_success_response(
)
api.pause_job_update(self.UPDATE_KEY, message='hello')
mock_proxy.pauseJobUpdate.assert_called_once_with(
self.UPDATE_KEY, 'hello')
def test_resume_job_update(self):
"""Test successful job update resume."""
api, mock_proxy = self.mock_api()
mock_proxy.resumeJobUpdate.return_value = self.create_simple_success_response(
)
api.resume_job_update(self.UPDATE_KEY, message='hello')
mock_proxy.resumeJobUpdate.assert_called_once_with(
self.UPDATE_KEY, 'hello')
def test_query_job_updates(self):
"""Test querying job updates."""
api, mock_proxy = self.mock_api()
job_key = AuroraJobKey("foo", "role", "env", "name")
query = JobUpdateQuery(
jobKey=job_key.to_thrift(),
updateStatuses={JobUpdateStatus.ROLLING_FORWARD})
api.query_job_updates(job_key=job_key,
update_statuses=query.updateStatuses)
mock_proxy.getJobUpdateSummaries.assert_called_once_with(query,
retry=True)
def test_query_job_updates_no_filter(self):
"""Test querying job updates with no filter args."""
api, mock_proxy = self.mock_api()
query = JobUpdateQuery()
api.query_job_updates()
mock_proxy.getJobUpdateSummaries.assert_called_once_with(query,
retry=True)
def test_get_job_update_details(self):
"""Test getting job update details."""
api, mock_proxy = self.mock_api()
key = JobUpdateKey(job=JobKey(role="role",
environment="env",
name="name"),
id="id")
api.get_job_update_details(key)
query = JobUpdateQuery(key=key)
mock_proxy.getJobUpdateDetails.assert_called_once_with(key,
query,
retry=True)
def test_set_quota(self):
"""Test setting quota."""
api, mock_proxy = self.mock_api()
api.set_quota("role", 1.0, 32, 64)
actual = list(mock_proxy.setQuota.mock_calls[0][1][1].resources)
assert Resource(numCpus=1.0) in actual
assert Resource(ramMb=32) in actual
assert Resource(diskMb=64) in actual
def test_update_status_json(self):
mock_context = FakeAuroraCommandContext()
api = mock_context.get_api('west')
api.query_job_updates.return_value = self.get_status_query_response()
api.get_job_update_details.return_value = self.get_update_details_response(
)
with contextlib.nested(
patch('apache.aurora.client.cli.update.Update.create_context',
return_value=mock_context),
patch('apache.aurora.client.factory.CLUSTERS',
new=self.TEST_CLUSTERS)):
cmd = AuroraCommandLine()
result = cmd.execute([
"beta-update", "status", "--write-json", "west/mcc/test/hello"
])
assert result == EXIT_OK
mock_context.get_api("west").query_job_updates.assert_called_with(
job_key=AuroraJobKey('west', 'mcc', 'test', 'hello'))
mock_context.get_api(
"west").get_job_update_details.assert_called_with('hello')
assert mock_context.get_out_str() == textwrap.dedent("""\
{
"status": "ROLLING_FORWARD",
"last_updated": 14114056030,
"started": 1411404927,
"update_events": [
{
"status": "ROLLING_FORWARD",
"timestampMs": 1411404927
},
{
"status": "ROLL_FORWARD_PAUSED",
"timestampMs": 1411405000
},
{
"status": "ROLLING_FORWARD",
"timestampMs": 1411405100
}
],
"job": "west/mcc/test/hello",
"updateId": "fake-update-identifier",
"instance_update_events": [
{
"action": "INSTANCE_UPDATING",
"instance": 1,
"timestamp": 1411404930
},
{
"action": "INSTANCE_UPDATING",
"instance": 2,
"timestamp": 1411404940
},
{
"action": "INSTANCE_UPDATED",
"instance": 1,
"timestamp": 1411404950
},
{
"action": "INSTANCE_UPDATED",
"instance": 2,
"timestamp": 1411404960
}
]
}""")
def job_key(self): return AuroraJobKey(self.cluster(), self.role(), self.environment(), self.name())
def execute(self, context):
update_filter = context.options.filter
cluster = update_filter.cluster
if (update_filter.role is not None and update_filter.env is not None
and update_filter.job is not None):
job_key = AuroraJobKey(cluster=cluster,
role=update_filter.role,
env=update_filter.env,
name=update_filter.job)
else:
job_key = None
api = context.get_api(cluster)
filter_statuses = set()
for status in context.options.status:
filter_statuses = filter_statuses.union(
set(self.STATUS_GROUPS[status]))
response = api.query_job_updates(
role=update_filter.role if job_key is None else None,
job_key=job_key,
update_statuses=filter_statuses if filter_statuses else None,
user=context.options.user)
context.log_response_and_raise(response)
# The API does not offer a way to query by environment, so if that filter is requested, we
# perform a more broad role-based query and filter here.
summaries = response.result.getJobUpdateSummariesResult.updateSummaries
if job_key is None and update_filter.env is not None:
summaries = [
s for s in summaries
if s.key.job.environment == update_filter.env
]
if context.options.write_json:
result = []
for summary in summaries:
job_entry = {
"job":
AuroraJobKey.from_thrift(cluster,
summary.key.job).to_path(),
"id":
summary.key.id,
"user":
summary.user,
"started":
format_timestamp(summary.state.createdTimestampMs),
"last_modified":
format_timestamp(summary.state.lastModifiedTimestampMs),
"status":
JobUpdateStatus._VALUES_TO_NAMES[summary.state.status]
}
result.append(job_entry)
context.print_out(
json.dumps(result,
indent=2,
separators=[',', ': '],
sort_keys=False))
else:
if summaries:
context.print_out(self.HEADER)
for summary in summaries:
context.print_out(
self.FORMAT_STR.format(
AuroraJobKey.from_thrift(cluster,
summary.key.job).to_path(),
summary.key.id,
JobUpdateStatus._VALUES_TO_NAMES[summary.state.status],
summary.user,
format_timestamp(summary.state.createdTimestampMs),
format_timestamp(
summary.state.lastModifiedTimestampMs)))
return EXIT_OK
def setUp(self):
self._scheduler = create_autospec(spec=SchedulerThriftApiSpec,
instance=True)
self._job_key = AuroraJobKey('cl', 'johndoe', 'test', 'test_job')
self._event = FakeEvent()
class TestJobUpdateApis(unittest.TestCase):
"""Job update APIs tests."""
UPDATE_CONFIG = {
'batch_size': 1,
'restart_threshold': 50,
'watch_secs': 50,
'max_per_shard_failures': 2,
'max_total_failures': 1,
'rollback_on_failure': True,
'wait_for_batch_completion': False,
}
JOB_KEY = AuroraJobKey("foo", "role", "env", "name")
@classmethod
def create_blank_response(cls, code, msg):
response = Mock(spec=Response)
response.responseCode = code
response.messageDEPRECATED = msg
response.result = Mock(spec=Result)
return response
@classmethod
def create_simple_success_response(cls):
return cls.create_blank_response(ResponseCode.OK, 'OK')
@classmethod
def create_error_response(cls):
return cls.create_blank_response(ResponseCode.ERROR, 'ERROR')
@classmethod
def mock_api(cls):
api = AuroraClientAPI(Cluster(name="foo"))
mock_proxy = Mock()
api._scheduler_proxy = mock_proxy
return api, mock_proxy
@classmethod
def create_update_settings(cls):
return JobUpdateSettings(updateGroupSize=1,
maxPerInstanceFailures=2,
maxFailedInstances=1,
maxWaitToInstanceRunningMs=50 * 1000,
minWaitInInstanceRunningMs=50 * 1000,
rollbackOnFailure=True)
@classmethod
def create_update_request(cls, task_config):
return JobUpdateRequest(jobKey=cls.JOB_KEY.to_thrift(),
instanceCount=5,
settings=cls.create_update_settings(),
taskConfig=task_config)
@classmethod
def mock_job_config(cls, error=None):
config = Mock(spec=AuroraConfig)
mock_get = Mock()
mock_get.get.return_value = cls.UPDATE_CONFIG
if error:
config.update_config.side_effect = error
else:
config.update_config.return_value = mock_get
mock_task_config = Mock()
mock_task_config.taskConfig = TaskConfig()
config.job.return_value = mock_task_config
config.role.return_value = "role"
config.environment.return_value = "env"
config.name.return_value = "name"
config.instances.return_value = 5
return config
def test_start_job_update(self):
"""Test successful job update start."""
api, mock_proxy = self.mock_api()
task_config = TaskConfig()
mock_proxy.startJobUpdate.return_value = self.create_simple_success_response(
)
api.start_job_update(self.mock_job_config())
mock_proxy.startJobUpdate.assert_called_once_with(
self.create_update_request(task_config))
def test_start_job_update_fails_parse_update_config(self):
"""Test start_job_update fails to parse invalid UpdateConfig."""
api, mock_proxy = self.mock_api()
self.assertRaises(AuroraClientAPI.UpdateConfigError,
api.start_job_update,
self.mock_job_config(error=ValueError()))
def test_pause_job_update(self):
"""Test successful job update pause."""
api, mock_proxy = self.mock_api()
mock_proxy.pauseJobUpdate.return_value = self.create_simple_success_response(
)
api.pause_job_update(self.JOB_KEY)
mock_proxy.pauseJobUpdate.assert_called_once_with(
self.JOB_KEY.to_thrift())
def test_pause_job_update_invalid_key(self):
"""Test job update pause with invalid job key."""
api, mock_proxy = self.mock_api()
self.assertRaises(AuroraClientAPI.TypeError, api.pause_job_update,
"invalid")
def test_resume_job_update(self):
"""Test successful job update resume."""
api, mock_proxy = self.mock_api()
mock_proxy.resumeJobUpdate.return_value = self.create_simple_success_response(
)
api.resume_job_update(self.JOB_KEY)
mock_proxy.resumeJobUpdate.assert_called_once_with(
self.JOB_KEY.to_thrift())
def test_resume_job_update_invalid_key(self):
"""Test job update resume with invalid job key."""
api, mock_proxy = self.mock_api()
self.assertRaises(AuroraClientAPI.TypeError, api.resume_job_update,
"invalid")
def test_abort_job_update(self):
"""Test successful job update abort."""
api, mock_proxy = self.mock_api()
mock_proxy.abortJobUpdate.return_value = self.create_simple_success_response(
)
api.abort_job_update(self.JOB_KEY)
mock_proxy.abortJobUpdate.assert_called_once_with(
self.JOB_KEY.to_thrift())
def test_abort_job_update_invalid_key(self):
"""Test job update abort with invalid job key."""
api, mock_proxy = self.mock_api()
self.assertRaises(AuroraClientAPI.TypeError, api.abort_job_update,
"invalid")
def test_query_job_updates(self):
"""Test querying job updates."""
api, mock_proxy = self.mock_api()
job_key = AuroraJobKey("foo", "role", "env", "name")
query = JobUpdateQuery(jobKey=job_key.to_thrift(),
updateStatuses=set(
[JobUpdateStatus.ROLLING_FORWARD]))
api.query_job_updates(job_key=job_key,
update_statuses=query.updateStatuses)
mock_proxy.getJobUpdateSummaries.assert_called_once_with(query)
def test_query_job_updates_no_filter(self):
"""Test querying job updates with no filter args."""
api, mock_proxy = self.mock_api()
query = JobUpdateQuery()
api.query_job_updates()
mock_proxy.getJobUpdateSummaries.assert_called_once_with(query)
def test_get_job_update_details(self):
"""Test getting job update details."""
api, mock_proxy = self.mock_api()
api.get_job_update_details("id")
mock_proxy.getJobUpdateDetails.assert_called_once_with("id")
