Repository: aurora Updated Branches: refs/heads/master d521dcd72 -> 32a8a0772
AURORA-1492 Improve "aurora update start" command output Reviewed at https://reviews.apache.org/r/47550/ Project: http://git-wip-us.apache.org/repos/asf/aurora/repo Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/32a8a077 Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/32a8a077 Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/32a8a077 Branch: refs/heads/master Commit: 32a8a077293c53ffe6f38bd35a6030984d049220 Parents: d521dcd Author: Mehrdad Nurolahzade <[email protected]> Authored: Wed May 25 17:39:49 2016 -0500 Committer: Joshua Cohen <[email protected]> Committed: Wed May 25 17:39:49 2016 -0500 ---------------------------------------------------------------------- .../apache/aurora/client/cli/diff_formatter.py | 139 ++++++++++++++ .../python/apache/aurora/client/cli/jobs.py | 118 +----------- .../python/apache/aurora/client/cli/update.py | 3 + .../apache/aurora/client/cli/test_diff.py | 138 +++---------- .../aurora/client/cli/test_diff_formatter.py | 192 +++++++++++++++++++ .../apache/aurora/client/cli/test_supdate.py | 67 +++++-- 6 files changed, 417 insertions(+), 240 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/main/python/apache/aurora/client/cli/diff_formatter.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/cli/diff_formatter.py b/src/main/python/apache/aurora/client/cli/diff_formatter.py new file mode 100644 index 0000000..a1b9bb3 --- /dev/null +++ b/src/main/python/apache/aurora/client/cli/diff_formatter.py @@ -0,0 +1,139 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import os +import pprint +import subprocess +from copy import deepcopy +from itertools import chain +from pipes import quote +from tempfile import NamedTemporaryFile + +from apache.aurora.client.cli import EXIT_COMMAND_FAILURE, EXIT_INVALID_PARAMETER + +from gen.apache.aurora.api.constants import ACTIVE_STATES, AURORA_EXECUTOR_NAME +from gen.apache.aurora.api.ttypes import ExecutorConfig + + +class DiffFormatter(object): + def __init__(self, context, config, cluster=None, role=None, env=None, name=None): + self.context = context + self.config = config + self.cluster = config.cluster() if cluster is None else cluster + self.role = config.role() if role is None else role + self.env = config.environment() if env is None else env + self.name = config.name() if name is None else name + self.prettyprinter = pprint.PrettyPrinter(indent=2) + + def _pretty_print_task(self, task): + task.configuration = None + executor_config = json.loads(task.executorConfig.data) + pretty_executor = json.dumps(executor_config, indent=2, sort_keys=True) + pretty_executor = '\n'.join([" %s" % s for s in pretty_executor.split("\n")]) + # Make start cleaner, and display multi-line commands across multiple lines. + pretty_executor = '\n ' + pretty_executor.replace(r'\n', '\n') + + # Avoid re-escaping as it's already pretty printed. + class RawRepr(object): + def __init__(self, data): + self.data = data + + def __repr__(self): + return self.data + + task.executorConfig = ExecutorConfig( + name=AURORA_EXECUTOR_NAME, + data=RawRepr(pretty_executor)) + return self.prettyprinter.pformat(vars(task)) + + def _pretty_print_tasks(self, tasks): + return ",\n".join(self._pretty_print_task(t) for t in tasks) + + def _dump_tasks(self, tasks, out_file): + out_file.write(self._pretty_print_tasks(tasks)) + out_file.write("\n") + out_file.flush() + + def _diff_tasks(self, local_tasks, remote_tasks): + diff_program = os.environ.get("DIFF_VIEWER", "diff") + with NamedTemporaryFile() as local: + self._dump_tasks(local_tasks, local) + with NamedTemporaryFile() as remote: + self._dump_tasks(remote_tasks, remote) + result = subprocess.call("%s %s %s" % ( + diff_program, quote(remote.name), quote(local.name)), shell=True) + # Unlike most commands, diff doesn't return zero on success; it returns + # 1 when a successful diff is non-empty. + if result not in (0, 1): + raise self.context.CommandError(EXIT_COMMAND_FAILURE, "Error running diff command") + + def _show_diff(self, header, configs_summaries, local_task=None): + def min_start(ranges): + return min(ranges, key=lambda r: r.first).first + + def format_ranges(ranges): + instances = [] + for task_range in sorted(list(ranges), key=lambda r: r.first): + if task_range.first == task_range.last: + instances.append("[%s]" % task_range.first) + else: + instances.append("[%s-%s]" % (task_range.first, task_range.last)) + return instances + + def print_instances(instances): + self.context.print_out("%s %s" % (header, ", ".join(str(span) for span in instances))) + + summaries = sorted(list(configs_summaries), key=lambda s: min_start(s.instances)) + + if local_task: + for summary in summaries: + print_instances(format_ranges(summary.instances)) + self.context.print_out("with diff:\n") + self._diff_tasks([deepcopy(local_task)], [summary.config]) + self.context.print_out('') + else: + if summaries: + print_instances( + format_ranges(r for r in chain.from_iterable(s.instances for s in summaries))) + + def diff_no_update_details(self, local_tasks): + api = self.context.get_api(self.cluster) + resp = api.query(api.build_query(self.role, self.name, env=self.env, statuses=ACTIVE_STATES)) + self.context.log_response_and_raise( + resp, + err_code=EXIT_INVALID_PARAMETER, + err_msg="Could not find job to diff against") + if resp.result.scheduleStatusResult.tasks is None: + self.context.print_err("No tasks found for job %s" % + self.context.options.instance_spec.jobkey) + return EXIT_COMMAND_FAILURE + else: + remote_tasks = [t.assignedTask.task for t in resp.result.scheduleStatusResult.tasks] + self._diff_tasks(local_tasks, remote_tasks) + + def show_job_update_diff(self, instances, local_task=None): + api = self.context.get_api(self.cluster) + resp = api.get_job_update_diff(self.config, instances) + self.context.log_response_and_raise( + resp, + err_code=EXIT_COMMAND_FAILURE, + err_msg="Error getting diff info from scheduler") + diff = resp.result.getJobUpdateDiffResult + self.context.print_out("This job update will:") + self._show_diff("add instances:", diff.add) + self._show_diff("remove instances:", diff.remove) + self._show_diff("remove instances:", diff.remove) + self._show_diff("update instances:", diff.update, local_task) + self._show_diff("not change instances:", diff.unchanged) http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/main/python/apache/aurora/client/cli/jobs.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/cli/jobs.py b/src/main/python/apache/aurora/client/cli/jobs.py index 336d6fa..7b4c269 100644 --- a/src/main/python/apache/aurora/client/cli/jobs.py +++ b/src/main/python/apache/aurora/client/cli/jobs.py @@ -16,24 +16,18 @@ from __future__ import print_function import json import logging -import os import pprint -import subprocess import textwrap import webbrowser from collections import namedtuple from copy import deepcopy from datetime import datetime -from itertools import chain -from pipes import quote -from tempfile import NamedTemporaryFile from thrift.protocol import TJSONProtocol from thrift.TSerialization import serialize from apache.aurora.client.api.job_monitor import JobMonitor from apache.aurora.client.api.restarter import RestartSettings -from apache.aurora.client.api.scheduler_client import SchedulerProxy from apache.aurora.client.base import get_job_page, synthesize_url from apache.aurora.client.cli import ( EXIT_COMMAND_FAILURE, @@ -45,6 +39,7 @@ from apache.aurora.client.cli import ( Verb ) from apache.aurora.client.cli.context import AuroraCommandContext +from apache.aurora.client.cli.diff_formatter import DiffFormatter from apache.aurora.client.cli.options import ( ADD_INSTANCE_WAIT_OPTION, ALL_INSTANCES, @@ -69,8 +64,8 @@ from apache.aurora.client.cli.options import ( from apache.aurora.common.aurora_job_key import AuroraJobKey from apache.aurora.config.resource import ResourceManager -from gen.apache.aurora.api.constants import ACTIVE_STATES, AURORA_EXECUTOR_NAME -from gen.apache.aurora.api.ttypes import ExecutorConfig, ResponseCode, ScheduleStatus +from gen.apache.aurora.api.constants import ACTIVE_STATES +from gen.apache.aurora.api.ttypes import ResponseCode, ScheduleStatus # Utility type, representing job keys with wildcards. PartialJobKey = namedtuple('PartialJobKey', ['cluster', 'role', 'env', 'name']) @@ -184,76 +179,6 @@ class DiffCommand(Verb): INSTANCES_SPEC_ARGUMENT, CONFIG_ARGUMENT] - def pretty_print_task(self, task): - task.configuration = None - executor_config = json.loads(task.executorConfig.data) - pretty_executor = json.dumps(executor_config, indent=2, sort_keys=True) - pretty_executor = '\n'.join([" %s" % s for s in pretty_executor.split("\n")]) - # Make start cleaner, and display multi-line commands across multiple lines. - pretty_executor = '\n ' + pretty_executor.replace(r'\n', '\n') - # Avoid re-escaping as it's already pretty printed. - class RawRepr(object): - def __init__(self, data): - self.data = data - - def __repr__(self): - return self.data - - task.executorConfig = ExecutorConfig( - name=AURORA_EXECUTOR_NAME, - data=RawRepr(pretty_executor)) - return self.prettyprinter.pformat(vars(task)) - - def pretty_print_tasks(self, tasks): - return ",\n".join(self.pretty_print_task(t) for t in tasks) - - def dump_tasks(self, tasks, out_file): - out_file.write(self.pretty_print_tasks(tasks)) - out_file.write("\n") - out_file.flush() - - def diff_tasks(self, context, local_tasks, remote_tasks): - diff_program = os.environ.get("DIFF_VIEWER", "diff") - with NamedTemporaryFile() as local: - self.dump_tasks(local_tasks, local) - with NamedTemporaryFile() as remote: - self.dump_tasks(remote_tasks, remote) - result = subprocess.call("%s %s %s" % ( - diff_program, quote(remote.name), quote(local.name)), shell=True) - # Unlike most commands, diff doesn't return zero on success; it returns - # 1 when a successful diff is non-empty. - if result not in (0, 1): - raise context.CommandError(EXIT_COMMAND_FAILURE, "Error running diff command") - - def show_diff(self, context, header, configs_summaries, local_task=None): - def min_start(ranges): - return min(ranges, key=lambda r: r.first).first - - def format_ranges(ranges): - instances = [] - for task_range in sorted(list(ranges), key=lambda r: r.first): - if task_range.first == task_range.last: - instances.append("[%s]" % task_range.first) - else: - instances.append("[%s-%s]" % (task_range.first, task_range.last)) - return instances - - def print_instances(instances): - context.print_out("%s %s" % (header, ", ".join(str(span) for span in instances))) - - summaries = sorted(list(configs_summaries), key=lambda s: min_start(s.instances)) - - if local_task: - for summary in summaries: - print_instances(format_ranges(summary.instances)) - context.print_out("with diff:\n") - self.diff_tasks(context, [deepcopy(local_task)], [summary.config]) - context.print_out('') - else: - if summaries: - print_instances( - format_ranges(r for r in chain.from_iterable(s.instances for s in summaries))) - def execute(self, context): config = context.get_job_config( context.options.instance_spec.jobkey, @@ -280,41 +205,14 @@ class DiffCommand(Verb): local_tasks = [ deepcopy(local_task) for _ in range(config.instances()) ] - - def diff_no_update_details(): - resp = api.query(api.build_query(role, name, env=env, statuses=ACTIVE_STATES)) - context.log_response_and_raise( - resp, - err_code=EXIT_INVALID_PARAMETER, - err_msg="Could not find job to diff against") - if resp.result.scheduleStatusResult.tasks is None: - context.print_err("No tasks found for job %s" % context.options.instance_spec.jobkey) - return EXIT_COMMAND_FAILURE - else: - remote_tasks = [t.assignedTask.task for t in resp.result.scheduleStatusResult.tasks] - self.diff_tasks(context, local_tasks, remote_tasks) + instances = (None if context.options.instance_spec.instance == ALL_INSTANCES else + context.options.instance_spec.instance) + formatter = DiffFormatter(context, config, cluster, role, env, name) if config.raw().has_cron_schedule(): - diff_no_update_details() + formatter.diff_no_update_details(local_tasks) else: - instances = (None if context.options.instance_spec.instance == ALL_INSTANCES else - context.options.instance_spec.instance) - try: - resp = api.get_job_update_diff(config, instances) - context.log_response_and_raise( - resp, - err_code=EXIT_COMMAND_FAILURE, - err_msg="Error getting diff info from scheduler") - diff = resp.result.getJobUpdateDiffResult - context.print_out("This job update will:") - self.show_diff(context, "add instances:", diff.add) - self.show_diff(context, "remove instances:", diff.remove) - self.show_diff(context, "update instances:", diff.update, local_task) - self.show_diff(context, "not change instances:", diff.unchanged) - except SchedulerProxy.ThriftInternalError: - # TODO(maxim): Temporary fallback to ensure client/scheduler backwards compatibility - # (i.e. new client works against old scheduler). - diff_no_update_details() + formatter.show_job_update_diff(instances, local_task) return EXIT_OK http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/main/python/apache/aurora/client/cli/update.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/cli/update.py b/src/main/python/apache/aurora/client/cli/update.py index eb7e9b0..bb526f7 100644 --- a/src/main/python/apache/aurora/client/cli/update.py +++ b/src/main/python/apache/aurora/client/cli/update.py @@ -33,6 +33,7 @@ from apache.aurora.client.cli import ( Verb ) from apache.aurora.client.cli.context import AuroraCommandContext +from apache.aurora.client.cli.diff_formatter import DiffFormatter from apache.aurora.client.cli.options import ( ALL_INSTANCES, BIND_OPTION, @@ -169,6 +170,8 @@ class StartUpdate(Verb): "Cron jobs may only be updated with \"aurora cron schedule\" command") api = context.get_api(config.cluster()) + formatter = DiffFormatter(context, config) + formatter.show_job_update_diff(instances) try: resp = api.start_job_update(config, context.options.message, instances) except AuroraClientAPI.UpdateConfigError as e: http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/test/python/apache/aurora/client/cli/test_diff.py ---------------------------------------------------------------------- diff --git a/src/test/python/apache/aurora/client/cli/test_diff.py b/src/test/python/apache/aurora/client/cli/test_diff.py index b9e91cf..d167bc4 100644 --- a/src/test/python/apache/aurora/client/cli/test_diff.py +++ b/src/test/python/apache/aurora/client/cli/test_diff.py @@ -12,15 +12,11 @@ # limitations under the License. # -import contextlib -import os -import textwrap - from mock import Mock, call, patch from pystachio import Empty -from apache.aurora.client.api import SchedulerProxy from apache.aurora.client.cli import EXIT_OK +from apache.aurora.client.cli.diff_formatter import DiffFormatter from apache.aurora.client.cli.jobs import DiffCommand from apache.aurora.client.cli.options import TaskInstanceKey from apache.aurora.config import AuroraConfig @@ -29,17 +25,7 @@ from apache.thermos.config.schema_base import MB, Process, Resources, Task from .util import AuroraClientCommandTest, FakeAuroraCommandContext, mock_verb_options -from gen.apache.aurora.api.constants import ACTIVE_STATES -from gen.apache.aurora.api.ttypes import ( - ConfigGroup, - GetJobUpdateDiffResult, - PopulateJobResult, - Range, - ResponseCode, - Result, - ScheduleStatusResult, - TaskQuery -) +from gen.apache.aurora.api.ttypes import PopulateJobResult, Result class TestDiffCommand(AuroraClientCommandTest): @@ -50,6 +36,7 @@ class TestDiffCommand(AuroraClientCommandTest): self._fake_context = FakeAuroraCommandContext() self._fake_context.set_options(self._mock_options) self._mock_api = self._fake_context.get_api("test") + self._formatter = Mock(spec=DiffFormatter) @classmethod def get_job_config(self, is_cron=False): @@ -69,113 +56,42 @@ class TestDiffCommand(AuroraClientCommandTest): )) @classmethod - def create_status_response(cls): - resp = cls.create_simple_success_response() - resp.result = Result( - scheduleStatusResult=ScheduleStatusResult(tasks=set(cls.create_scheduled_tasks()))) - return resp - - @classmethod - def create_failed_status_response(cls): - return cls.create_blank_response(ResponseCode.INVALID_REQUEST, 'No tasks found for query') - - @classmethod def populate_job_config_result(cls): populate = cls.create_simple_success_response() populate.result = Result(populateJobResult=PopulateJobResult( taskConfig=cls.create_scheduled_tasks()[0].assignedTask.task)) return populate - @classmethod - def get_job_update_diff_result(cls): - diff = cls.create_simple_success_response() - task = cls.create_task_config('foo') - diff.result = Result(getJobUpdateDiffResult=GetJobUpdateDiffResult( - add=set([ConfigGroup( - config=task, - instances=frozenset([Range(first=10, last=10), Range(first=12, last=14)]))]), - remove=frozenset(), - update=frozenset([ConfigGroup( - config=task, - instances=frozenset([Range(first=11, last=11)]))]), - unchanged=frozenset([ConfigGroup( - config=task, - instances=frozenset([Range(first=0, last=9)]))]) - )) - return diff - def test_service_diff(self): config = self.get_job_config() self._fake_context.get_job_config = Mock(return_value=config) - self._mock_api.populate_job_config.return_value = self.populate_job_config_result() - self._mock_api.get_job_update_diff.return_value = self.get_job_update_diff_result() - - with contextlib.nested( - patch('subprocess.call', return_value=0), - patch('json.loads', return_value={})) as (subprocess_patch, _): + resp = self.populate_job_config_result() + self._mock_api.populate_job_config.return_value = resp - result = self._command.execute(self._fake_context) + with patch('apache.aurora.client.cli.jobs.DiffFormatter') as formatter: + formatter.return_value = self._formatter + assert self._command.execute(self._fake_context) == EXIT_OK - assert result == EXIT_OK - assert self._mock_api.populate_job_config.mock_calls == [call(config)] - assert self._mock_api.get_job_update_diff.mock_calls == [ - call(config, self._mock_options.instance_spec.instance) - ] - assert "\n".join(self._fake_context.get_out()) == textwrap.dedent("""\ - This job update will: - add instances: [10], [12-14] - update instances: [11] - with diff:\n\n - not change instances: [0-9]""") - assert subprocess_patch.call_count == 1 - assert subprocess_patch.call_args[0][0].startswith( - os.environ.get('DIFF_VIEWER', 'diff') + ' ') - - def test_service_diff_old_api(self): - config = self.get_job_config() - query = TaskQuery( - jobKeys=[self.TEST_JOBKEY.to_thrift()], - statuses=ACTIVE_STATES) - self._fake_context.get_job_config = Mock(return_value=config) - self._mock_api.populate_job_config.return_value = self.populate_job_config_result() - self._mock_api.get_job_update_diff.side_effect = SchedulerProxy.ThriftInternalError("Expected") - self._mock_api.query.return_value = self.create_empty_task_result() - self._mock_api.build_query.return_value = query - - with contextlib.nested( - patch('subprocess.call', return_value=0), - patch('json.loads', return_value={})) as (subprocess_patch, _): - - result = self._command.execute(self._fake_context) - assert result == EXIT_OK - assert self._mock_api.populate_job_config.mock_calls == [call(config)] - assert self._mock_api.get_job_update_diff.mock_calls == [ - call(config, self._mock_options.instance_spec.instance) - ] - assert self._mock_api.query.mock_calls == [call(query)] - assert subprocess_patch.call_count == 1 - assert subprocess_patch.call_args[0][0].startswith( - os.environ.get('DIFF_VIEWER', 'diff') + ' ') + assert self._mock_api.populate_job_config.mock_calls == [call(config)] + assert self._formatter.show_job_update_diff.mock_calls == [ + call(self._mock_options.instance_spec.instance, resp.result.populateJobResult.taskConfig) + ] + assert self._fake_context.get_out() == [] + assert self._fake_context.get_err() == [] def test_cron_diff(self): config = self.get_job_config(is_cron=True) - query = TaskQuery( - jobKeys=[self.TEST_JOBKEY.to_thrift()], - statuses=ACTIVE_STATES) self._fake_context.get_job_config = Mock(return_value=config) - self._mock_api.populate_job_config.return_value = self.populate_job_config_result() - self._mock_api.query.return_value = self.create_empty_task_result() - self._mock_api.build_query.return_value = query - - with contextlib.nested( - patch('subprocess.call', return_value=0), - patch('json.loads', return_value={})) as (subprocess_patch, _): - - result = self._command.execute(self._fake_context) - - assert result == EXIT_OK - assert self._mock_api.populate_job_config.mock_calls == [call(config)] - assert self._mock_api.query.mock_calls == [call(query)] - assert subprocess_patch.call_count == 1 - assert subprocess_patch.call_args[0][0].startswith( - os.environ.get('DIFF_VIEWER', 'diff') + ' ') + resp = self.populate_job_config_result() + self._mock_api.populate_job_config.return_value = resp + + with patch('apache.aurora.client.cli.jobs.DiffFormatter') as formatter: + formatter.return_value = self._formatter + assert self._command.execute(self._fake_context) == EXIT_OK + + assert self._mock_api.populate_job_config.mock_calls == [call(config)] + assert self._formatter.diff_no_update_details.mock_calls == [ + call([resp.result.populateJobResult.taskConfig] * 3) + ] + assert self._fake_context.get_out() == [] + assert self._fake_context.get_err() == [] http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/test/python/apache/aurora/client/cli/test_diff_formatter.py ---------------------------------------------------------------------- diff --git a/src/test/python/apache/aurora/client/cli/test_diff_formatter.py b/src/test/python/apache/aurora/client/cli/test_diff_formatter.py new file mode 100644 index 0000000..a145d88 --- /dev/null +++ b/src/test/python/apache/aurora/client/cli/test_diff_formatter.py @@ -0,0 +1,192 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import contextlib +import os +import textwrap + +import pytest +from mock import Mock, call, patch +from pystachio import Empty + +from apache.aurora.client.cli import Context +from apache.aurora.client.cli.diff_formatter import DiffFormatter +from apache.aurora.client.cli.jobs import DiffCommand +from apache.aurora.client.cli.options import TaskInstanceKey +from apache.aurora.config import AuroraConfig +from apache.aurora.config.schema.base import Job +from apache.thermos.config.schema_base import MB, Process, Resources, Task + +from .util import AuroraClientCommandTest, FakeAuroraCommandContext, mock_verb_options + +from gen.apache.aurora.api.constants import ACTIVE_STATES +from gen.apache.aurora.api.ttypes import ( + ConfigGroup, + GetJobUpdateDiffResult, + Range, + Result, + TaskQuery +) + + +class TestDiffFormatter(AuroraClientCommandTest): + def setUp(self): + self._fake_context = FakeAuroraCommandContext() + self._mock_options = mock_verb_options(DiffCommand()) + self._mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [0, 1]) + self._fake_context.set_options(self._mock_options) + self._mock_api = self._fake_context.get_api("west") + + @classmethod + def get_job_config(self, is_cron=False): + return AuroraConfig(job=Job( + cluster='west', + role='bozo', + environment='test', + name='the_job', + service=True if not is_cron else False, + cron_schedule='* * * * *' if is_cron else Empty, + task=Task( + name='task', + processes=[Process(cmdline='ls -la', name='process')], + resources=Resources(cpu=1.0, ram=1024 * MB, disk=1024 * MB) + ), + instances=3, + )) + + @classmethod + def get_job_update_diff_result(cls): + diff = cls.create_simple_success_response() + task = cls.create_task_config('foo') + diff.result = Result(getJobUpdateDiffResult=GetJobUpdateDiffResult( + add=set([ConfigGroup( + config=task, + instances=frozenset([Range(first=10, last=10), Range(first=12, last=14)]))]), + remove=frozenset(), + update=frozenset([ConfigGroup( + config=task, + instances=frozenset([Range(first=11, last=11)]))]), + unchanged=frozenset([ConfigGroup( + config=task, + instances=frozenset([Range(first=0, last=9)]))]) + )) + return diff + + @classmethod + def get_job_update_no_change_diff_result(cls): + diff = cls.create_simple_success_response() + task = cls.create_task_config('foo') + diff.result = Result(getJobUpdateDiffResult=GetJobUpdateDiffResult( + add=frozenset(), + remove=frozenset(), + update=frozenset(), + unchanged=frozenset([ConfigGroup( + config=task, + instances=frozenset([Range(first=0, last=3)]))]) + )) + return diff + + def test_show_job_update_diff_with_task_diff(self): + config = self.get_job_config() + self._fake_context.get_job_config = Mock(return_value=config) + formatter = DiffFormatter(self._fake_context, config) + local_task = self.create_scheduled_tasks()[0].assignedTask.task + self._mock_api.get_job_update_diff.return_value = self.get_job_update_diff_result() + + with contextlib.nested( + patch('subprocess.call', return_value=0), + patch('json.loads', return_value={})) as (subprocess_patch, _): + + formatter.show_job_update_diff(self._mock_options.instance_spec.instance, local_task) + + assert self._mock_api.get_job_update_diff.mock_calls == [ + call(config, self._mock_options.instance_spec.instance) + ] + assert "\n".join(self._fake_context.get_out()) == textwrap.dedent("""\ + This job update will: + add instances: [10], [12-14] + update instances: [11] + with diff:\n\n + not change instances: [0-9]""") + assert subprocess_patch.call_count == 1 + assert subprocess_patch.call_args[0][0].startswith( + os.environ.get('DIFF_VIEWER', 'diff') + ' ') + + def test_show_job_update_diff_without_task_diff(self): + config = self.get_job_config() + self._fake_context.get_job_config = Mock(return_value=config) + formatter = DiffFormatter(self._fake_context, config) + self._mock_api.get_job_update_diff.return_value = self.get_job_update_diff_result() + + formatter.show_job_update_diff(self._mock_options.instance_spec.instance) + + assert self._mock_api.get_job_update_diff.mock_calls == [ + call(config, self._mock_options.instance_spec.instance) + ] + assert "\n".join(self._fake_context.get_out()) == textwrap.dedent("""\ + This job update will: + add instances: [10], [12-14] + update instances: [11] + not change instances: [0-9]""") + + def test_show_job_update_diff_no_change(self): + config = self.get_job_config() + self._fake_context.get_job_config = Mock(return_value=config) + formatter = DiffFormatter(self._fake_context, config) + self._mock_api.get_job_update_diff.return_value = self.get_job_update_no_change_diff_result() + + formatter.show_job_update_diff(self._mock_options.instance_spec.instance) + + assert self._mock_api.get_job_update_diff.mock_calls == [ + call(config, self._mock_options.instance_spec.instance) + ] + assert "\n".join(self._fake_context.get_out()) == textwrap.dedent("""\ + This job update will: + not change instances: [0-3]""") + + def test_get_job_update_diff_error(self): + mock_config = self.get_job_config() + self._fake_context.get_job_config = Mock(return_value=mock_config) + formatter = DiffFormatter(self._fake_context, mock_config) + self._mock_api.get_job_update_diff.return_value = self.create_error_response() + + with pytest.raises(Context.CommandError): + formatter.show_job_update_diff(self._mock_options.instance_spec.instance) + + assert self._mock_api.get_job_update_diff.mock_calls == [ + call(mock_config, self._mock_options.instance_spec.instance) + ] + assert self._fake_context.get_out() == [] + assert self._fake_context.get_err() == ["Error getting diff info from scheduler", "\tWhoops"] + + def test_diff_no_update_details_success(self): + config = self.get_job_config(is_cron=True) + self._fake_context.get_job_config = Mock(return_value=config) + formatter = DiffFormatter(self._fake_context, config) + self._mock_api.query.return_value = self.create_empty_task_result() + query = TaskQuery( + jobKeys=[self.TEST_JOBKEY.to_thrift()], + statuses=ACTIVE_STATES) + self._mock_api.build_query.return_value = query + local_tasks = [] + + with contextlib.nested( + patch('subprocess.call', return_value=0), + patch('json.loads', return_value={})) as (subprocess_patch, _): + + formatter.diff_no_update_details(local_tasks) + + assert subprocess_patch.call_count == 1 + assert subprocess_patch.call_args[0][0].startswith( + os.environ.get('DIFF_VIEWER', 'diff') + ' ') http://git-wip-us.apache.org/repos/asf/aurora/blob/32a8a077/src/test/python/apache/aurora/client/cli/test_supdate.py ---------------------------------------------------------------------- diff --git a/src/test/python/apache/aurora/client/cli/test_supdate.py b/src/test/python/apache/aurora/client/cli/test_supdate.py index 2135ca9..317b175 100644 --- a/src/test/python/apache/aurora/client/cli/test_supdate.py +++ b/src/test/python/apache/aurora/client/cli/test_supdate.py @@ -16,7 +16,7 @@ import time import mock import pytest -from mock import ANY, Mock, call, create_autospec +from mock import ANY, Mock, call, create_autospec, patch from pystachio import Empty from apache.aurora.client.cli import ( @@ -27,6 +27,7 @@ from apache.aurora.client.cli import ( EXIT_UNKNOWN_ERROR, Context ) +from apache.aurora.client.cli.diff_formatter import DiffFormatter from apache.aurora.client.cli.options import TaskInstanceKey from apache.aurora.client.cli.update import ( AbortUpdate, @@ -97,6 +98,7 @@ class TestStartUpdate(AuroraClientCommandTest): self._fake_context = FakeAuroraCommandContext() self._fake_context.set_options(self._mock_options) self._mock_api = self._fake_context.get_api('UNUSED') + self._formatter = Mock(spec=DiffFormatter) @classmethod def create_mock_config(cls, is_cron=False): @@ -114,8 +116,9 @@ class TestStartUpdate(AuroraClientCommandTest): ResponseCode.LOCK_ERROR, "Error.") - with pytest.raises(Context.CommandError): - self._command.execute(self._fake_context) + with patch('apache.aurora.client.cli.update.DiffFormatter'): + with pytest.raises(Context.CommandError): + self._command.execute(self._fake_context) assert self._mock_api.start_job_update.mock_calls == [ call(mock_config, None, self._mock_options.instance_spec.instance) @@ -137,8 +140,13 @@ class TestStartUpdate(AuroraClientCommandTest): self._fake_context.get_job_config = Mock(return_value=mock_config) self._mock_api.start_job_update.return_value = self.create_simple_success_response() - self._command.execute(self._fake_context) + with patch('apache.aurora.client.cli.update.DiffFormatter') as formatter: + formatter.return_value = self._formatter + self._command.execute(self._fake_context) + assert self._formatter.show_job_update_diff.mock_calls == [ + call(self._mock_options.instance_spec.instance) + ] assert self._mock_api.start_job_update.mock_calls == [ call(mock_config, None, self._mock_options.instance_spec.instance) ] @@ -152,14 +160,20 @@ class TestStartUpdate(AuroraClientCommandTest): self._fake_context.get_job_config = Mock(return_value=mock_config) self._mock_options.instance_spec = TaskInstanceKey(self._job_key, None) self._mock_options.message = 'hello' - assert self._command.execute(self._fake_context) == EXIT_OK + with patch('apache.aurora.client.cli.update.DiffFormatter') as formatter: + formatter.return_value = self._formatter + assert self._command.execute(self._fake_context) == EXIT_OK + + assert self._formatter.show_job_update_diff.mock_calls == [ + call(self._mock_options.instance_spec.instance) + ] assert self._mock_api.start_job_update.mock_calls == [ call(ANY, 'hello', None) ] assert self._fake_context.get_out() == [ - StartUpdate.UPDATE_MSG_TEMPLATE % - ('http://something_or_other/scheduler/role/env/name/update/id') + StartUpdate.UPDATE_MSG_TEMPLATE % + ('http://something_or_other/scheduler/role/env/name/update/id'), ] assert self._fake_context.get_err() == [] @@ -176,17 +190,22 @@ class TestStartUpdate(AuroraClientCommandTest): get_status_query_response(status=JobUpdateStatus.ROLLED_FORWARD) ] - assert self._command.execute(self._fake_context) == EXIT_OK + with patch('apache.aurora.client.cli.update.DiffFormatter') as formatter: + formatter.return_value = self._formatter + assert self._command.execute(self._fake_context) == EXIT_OK + assert self._formatter.show_job_update_diff.mock_calls == [ + call(self._mock_options.instance_spec.instance) + ] assert self._mock_api.start_job_update.mock_calls == [call(ANY, None, None)] assert self._mock_api.query_job_updates.mock_calls == [ - call(update_key=resp.result.startJobUpdateResult.key) + call(update_key=resp.result.startJobUpdateResult.key) ] assert self._fake_context.get_out() == [ - StartUpdate.UPDATE_MSG_TEMPLATE % - ('http://something_or_other/scheduler/role/env/name/update/id'), - 'Current state ROLLED_FORWARD' + StartUpdate.UPDATE_MSG_TEMPLATE % + ('http://something_or_other/scheduler/role/env/name/update/id'), + 'Current state ROLLED_FORWARD' ] assert self._fake_context.get_err() == [] @@ -203,7 +222,8 @@ class TestStartUpdate(AuroraClientCommandTest): get_status_query_response(status=JobUpdateStatus.ROLLED_BACK) ] - assert self._command.execute(self._fake_context) == EXIT_COMMAND_FAILURE + with patch('apache.aurora.client.cli.update.DiffFormatter'): + assert self._command.execute(self._fake_context) == EXIT_COMMAND_FAILURE def test_start_update_and_wait_error(self): mock_config = self.create_mock_config() @@ -218,7 +238,8 @@ class TestStartUpdate(AuroraClientCommandTest): get_status_query_response(status=JobUpdateStatus.ERROR) ] - assert self._command.execute(self._fake_context) == EXIT_UNKNOWN_ERROR + with patch('apache.aurora.client.cli.update.DiffFormatter'): + assert self._command.execute(self._fake_context) == EXIT_UNKNOWN_ERROR def test_start_update_command_line_succeeds_noop_update(self): resp = self.create_simple_success_response() @@ -226,11 +247,18 @@ class TestStartUpdate(AuroraClientCommandTest): self._mock_api.start_job_update.return_value = resp mock_config = self.create_mock_config() self._fake_context.get_job_config = Mock(return_value=mock_config) - result = self._command.execute(self._fake_context) - assert result == EXIT_OK + with patch('apache.aurora.client.cli.update.DiffFormatter') as formatter: + formatter.return_value = self._formatter + assert self._command.execute(self._fake_context) == EXIT_OK + + assert self._formatter.show_job_update_diff.mock_calls == [ + call(self._mock_options.instance_spec.instance) + ] assert self._mock_api.start_job_update.mock_calls == [call(ANY, None, None)] - assert self._fake_context.get_out() == ["Noop update."] + assert self._fake_context.get_out() == [ + "Noop update." + ] assert self._fake_context.get_err() == [] def test_update_pulse_interval_too_small(self): @@ -240,8 +268,9 @@ class TestStartUpdate(AuroraClientCommandTest): error = Context.CommandError(100, 'something failed') self._mock_api.start_job_update.side_effect = error - with pytest.raises(Context.CommandError) as e: - self._command.execute(self._fake_context) + with patch('apache.aurora.client.cli.update.DiffFormatter'): + with pytest.raises(Context.CommandError) as e: + self._command.execute(self._fake_context) assert e.value == error assert self._mock_api.start_job_update.mock_calls == [call(ANY, None, None)]
