def test_simple_successful_kill_job(self): """Run a test of the "kill" command against a mocked-out API: Verifies that the kill command sends the right API RPCs, and performs the correct tests on the result.""" mock_options = self.setup_mock_options() mock_config = Mock() mock_api_factory = self.setup_mock_api_factory() with contextlib.nested( patch('apache.aurora.client.commands.core.make_client_factory', return_value=mock_api_factory), patch('twitter.common.app.get_options', return_value=mock_options), patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_make_client_factory, options, mock_get_job_config): mock_api = mock_api_factory.return_value with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. self.assert_kill_job_called(mock_api) mock_api.kill_job.assert_called_with( AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV, name=self.TEST_JOB), None, config=mock_config) self.assert_scheduler_called(mock_api) assert mock_make_client_factory.call_count == 1
def test_kill_job_api_level_with_shards(self): """Test kill client-side API logic.""" mock_options = self.setup_mock_options() mock_options.shards = [0, 1, 2, 3] mock_config = Mock() mock_config.hooks = [] mock_config.raw.return_value.enable_hooks.return_value.get.return_value = False (mock_api, mock_scheduler_proxy) = self.create_mock_api() 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.factory.make_client', return_value=mock_api), patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS), patch('twitter.common.app.get_options', return_value=mock_options), patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)): with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. assert mock_scheduler_proxy.killTasks.call_count == 1 query = self.get_expected_task_query([0, 1, 2, 3]) mock_scheduler_proxy.killTasks.assert_called_with(query, None) self.assert_scheduler_called(mock_api, query, 3)
def test_kill_job_api_level_with_shards(self): """Test kill client-side API logic.""" mock_options = self.setup_mock_options() mock_options.shards = [0, 1, 2, 3] mock_config = Mock() mock_config.hooks = [] mock_config.raw.return_value.enable_hooks.return_value.get.return_value = False (mock_api, mock_scheduler) = self.setup_mock_api() mock_api_factory = Mock(return_value=mock_api) mock_scheduler.killTasks.return_value = self.get_kill_job_response() with contextlib.nested( patch('apache.aurora.client.factory.make_client_factory', return_value=mock_api_factory), patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler), patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS), patch('twitter.common.app.get_options', return_value=mock_options), patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_api_factory_patch, mock_scheduler_proxy_class, mock_clusters, options, mock_get_job_config): with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. self.assert_scheduler_called(mock_api) assert mock_scheduler.killTasks.call_count == 1 query = self.get_expected_task_query([0, 1, 2, 3]) mock_scheduler.killTasks.assert_called_with(query, None)
def test_kill_job_api_level_with_shards_batched(self): """Test kill client-side API logic.""" mock_options = self.setup_mock_options() mock_options.batch_size = 2 mock_options.shards = [0, 1, 2, 3] mock_config = Mock() mock_config.hooks = [] mock_config.raw.return_value.enable_hooks.return_value.get.return_value = False (mock_api, mock_scheduler_proxy) = self.create_mock_api() mock_api.check_status.return_value = self.create_status_call_result() mock_scheduler_proxy.getTasksWithoutConfigs.return_value = self.create_status_call_result() mock_scheduler_proxy.killTasks.return_value = self.get_kill_job_response() with contextlib.nested( patch('apache.aurora.client.factory.make_client', return_value=mock_api), patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS), patch('twitter.common.app.get_options', return_value=mock_options), patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config), patch('apache.aurora.client.commands.core.JobMonitor')): with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. assert mock_scheduler_proxy.killTasks.call_count == 2 query = self.get_expected_task_query([2, 3]) mock_scheduler_proxy.killTasks.assert_called_with(query, None)