| # |
| # 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 functools |
| |
| import pytest |
| from mock import call, create_autospec, patch |
| from twitter.common.contextutil import temporary_file |
| |
| from apache.aurora.client.api.health_check import StatusHealthCheck |
| from apache.aurora.client.cli import EXIT_API_ERROR, EXIT_INVALID_PARAMETER, Context |
| from apache.aurora.client.cli.client import AuroraCommandLine |
| from apache.aurora.client.cli.jobs import RestartCommand |
| from apache.aurora.client.cli.options import TaskInstanceKey |
| from apache.aurora.common.aurora_job_key import AuroraJobKey |
| |
| from .util import AuroraClientCommandTest, FakeAuroraCommandContext, IOMock, mock_verb_options |
| |
| from gen.apache.aurora.api.ttypes import JobKey, PopulateJobResult, Result, TaskConfig |
| |
| |
| class TestRestartJobCommand(AuroraClientCommandTest): |
| 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_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") |
| ] |
| |
| |
| class TestRestartCommand(AuroraClientCommandTest): |
| |
| @classmethod |
| def setup_mock_scheduler_for_simple_restart(cls, api): |
| """Set up all of the API mocks for scheduler calls during a simple restart""" |
| sched_proxy = api.scheduler_proxy |
| cls.setup_get_tasks_status_calls(sched_proxy) |
| cls.setup_populate_job_config(sched_proxy) |
| sched_proxy.restartShards.return_value = cls.create_simple_success_response() |
| |
| @classmethod |
| def setup_populate_job_config(cls, api): |
| populate = cls.create_simple_success_response() |
| populate.result = Result(populateJobResult=PopulateJobResult(taskConfig=TaskConfig())) |
| api.populateJobConfig.return_value = populate |
| return populate |
| |
| @classmethod |
| def setup_health_checks(cls): |
| mock_health_check = create_autospec(spec=StatusHealthCheck, instance=True) |
| mock_health_check.health.return_value = True |
| return mock_health_check |
| |
| def test_restart_simple(self): |
| # Test the client-side restart logic in its simplest case: everything succeeds |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait') |
| ): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello', |
| '--config', fp.name]) |
| |
| # Like the update test, the exact number of calls here doesn't matter. |
| # what matters is that it must have been called once before batching, plus |
| # at least once per batch, and there are 4 batches. |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count >= 4 |
| # called once per batch |
| assert mock_scheduler_proxy.restartShards.call_count == 4 |
| # parameters for all calls are generated by the same code, so we just check one |
| mock_scheduler_proxy.restartShards.assert_called_with(JobKey(environment=self.TEST_ENV, |
| role=self.TEST_ROLE, name=self.TEST_JOB), [15, 16, 17, 18, 19], retry=True) |
| |
| def test_restart_simple_no_config(self): |
| # Test the client-side restart logic in its simplest case: everything succeeds |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait') |
| ): |
| cmd = AuroraCommandLine() |
| cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello']) |
| |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count >= 4 |
| assert mock_scheduler_proxy.restartShards.call_count == 4 |
| mock_scheduler_proxy.restartShards.assert_called_with(JobKey(environment=self.TEST_ENV, |
| role=self.TEST_ROLE, name=self.TEST_JOB), [15, 16, 17, 18, 19], retry=True) |
| |
| def test_restart_invalid_shards(self): |
| # Test the client-side restart when a shard argument is too large, and it's |
| # using strict mode. |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait') |
| ): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| result = cmd.execute(['job', 'restart', '--batch-size=5', '--max-total-failures=-1', |
| 'west/bozo/test/hello', '--config', fp.name]) |
| assert result == EXIT_INVALID_PARAMETER |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 0 |
| assert mock_scheduler_proxy.restartShards.call_count == 0 |
| |
| def test_restart_failed_status(self): |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| mock_scheduler_proxy.getTasksWithoutConfigs.return_value = self.create_error_response() |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait')): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| result = cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello', |
| '--config', fp.name]) |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 1 |
| assert mock_scheduler_proxy.restartShards.call_count == 0 |
| assert result == EXIT_API_ERROR |
| |
| def test_restart_no_such_job_with_instances(self): |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| mock_io = IOMock() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| # Make getTasksWithoutConfigs return an error, which is what happens when a job is not found. |
| mock_scheduler_proxy.getTasksWithoutConfigs.return_value = self.create_error_response() |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.cli.context.AuroraCommandContext.print_err', |
| side_effect=mock_io.put), |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait')): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| result = cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello/1-3', |
| '--config', fp.name]) |
| # We need to check tat getTasksWithoutConfigs was called, but that restartShards wasn't. |
| # In older versions of the client, if shards were specified, but the job didn't |
| # exist, the error wouldn't be detected unti0 restartShards was called, which generated |
| # the wrong error message. |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 1 |
| assert mock_scheduler_proxy.restartShards.call_count == 0 |
| assert result == EXIT_API_ERROR |
| # Error message should be written to log, and it should be what was returned |
| # by the getTasksWithoutConfigs call. |
| assert mock_io.get() == ["Error restarting job west/bozo/test/hello:", "\tWhoops"] |
| |
| def test_restart_failed_restart(self): |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| mock_scheduler_proxy.restartShards.return_value = self.create_error_response() |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait')): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| result = cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello', |
| '--config', fp.name]) |
| assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 1 |
| assert mock_scheduler_proxy.restartShards.call_count == 1 |
| mock_scheduler_proxy.restartShards.assert_called_with(JobKey(environment=self.TEST_ENV, |
| role=self.TEST_ROLE, name=self.TEST_JOB), [0, 1, 2, 3, 4], retry=True) |
| assert result == EXIT_API_ERROR |
| |
| MOCK_OUT = [] |
| MOCK_ERR = [] |
| |
| @classmethod |
| def reset_mock_io(cls): |
| cls.MOCK_OUT = [] |
| cls.MOCK_ERR = [] |
| |
| @classmethod |
| def mock_print_out(cls, msg, indent=0): |
| indent_str = " " * indent |
| cls.MOCK_OUT.append("%s%s" % (indent_str, msg)) |
| |
| @classmethod |
| def mock_print_err(cls, msg, indent=0): |
| indent_str = " " * indent |
| cls.MOCK_ERR.append("%s%s" % (indent_str, msg)) |
| |
| def test_restart_simple_output(self): |
| self.reset_mock_io() |
| # Test the client-side restart logic in its simplest case: everything succeeds |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('threading._Event.wait'), |
| patch('apache.aurora.client.cli.context.AuroraCommandContext.print_out', |
| side_effect=self.mock_print_out), |
| patch('apache.aurora.client.cli.context.AuroraCommandContext.print_err', |
| side_effect=self.mock_print_err) |
| ): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello', |
| '--config', fp.name]) |
| assert self.MOCK_OUT == ['Job west/bozo/test/hello restarted successfully'] |
| assert self.MOCK_ERR == [] |
| |
| def test_restart_failed_restart_output(self): |
| self.reset_mock_io() |
| (mock_api, mock_scheduler_proxy) = self.create_mock_api() |
| mock_health_check = self.setup_health_checks() |
| self.setup_mock_scheduler_for_simple_restart(mock_api) |
| mock_scheduler_proxy.restartShards.return_value = self.create_error_response() |
| mock_scheduler_proxy.getTierConfigs.return_value = self.get_mock_tier_configurations() |
| with contextlib.nested( |
| patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), |
| patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck', |
| return_value=mock_health_check), |
| patch('time.time', side_effect=functools.partial(self.fake_time, self)), |
| patch('apache.aurora.client.cli.context.AuroraCommandContext.print_out', |
| side_effect=self.mock_print_out), |
| patch('apache.aurora.client.cli.context.AuroraCommandContext.print_err', |
| side_effect=self.mock_print_err), |
| patch('threading._Event.wait')): |
| with temporary_file() as fp: |
| fp.write(self.get_valid_config()) |
| fp.flush() |
| cmd = AuroraCommandLine() |
| cmd.execute(['job', 'restart', '--batch-size=5', 'west/bozo/test/hello', |
| '--config', fp.name]) |
| assert self.MOCK_OUT == [] |
| assert "Error restarting job west/bozo/test/hello:" in self.MOCK_ERR |