def test_simple_successful_kill_job(self): """Run a test of the "kill" command against a mocked-out API: Verifies that the kill command sends the right API RPCs, and performs the correct tests on the result.""" mock_options = self.setup_mock_options() mock_config = Mock() mock_api_factory = self.setup_mock_api_factory() with contextlib.nested( patch('twitter.aurora.client.commands.core.make_client_factory', return_value=mock_api_factory), patch('twitter.common.app.get_options', return_value=mock_options), patch('twitter.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_make_client_factory, options, mock_get_job_config): mock_api = mock_api_factory.return_value with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. self.assert_kill_job_called(mock_api) mock_api.kill_job.assert_called_with( AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV, name=self.TEST_JOB), None, config=mock_config) self.assert_scheduler_called(mock_api) assert mock_make_client_factory.call_count == 1
def test_simple_successful_kill_job(self): """Run a test of the "kill" command against a mocked-out API: Verifies that the kill command sends the right API RPCs, and performs the correct tests on the result.""" mock_options = self.setup_mock_options() mock_config = Mock() mock_api_factory = self.setup_mock_api_factory() with contextlib.nested( patch('twitter.aurora.client.commands.core.make_client_factory', return_value=mock_api_factory), patch('twitter.common.app.get_options', return_value=mock_options), patch('twitter.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_make_client_factory, options, mock_get_job_config): mock_api = mock_api_factory.return_value with temporary_file() as fp: fp.write(self.get_valid_config()) fp.flush() kill(['west/mchucarroll/test/hello', fp.name], mock_options) # Now check that the right API calls got made. self.assert_kill_job_called(mock_api) mock_api.kill_job.assert_called_with( AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV, name=self.TEST_JOB), None, config=mock_config) self.assert_scheduler_called(mock_api) assert mock_make_client_factory.call_count == 1
def 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('twitter.aurora.client.factory.make_client_factory', return_value=mock_api_factory), patch('twitter.aurora.client.api.SchedulerProxy', return_value=mock_scheduler), patch('twitter.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS), patch('twitter.common.app.get_options', return_value=mock_options), patch('twitter.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_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(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('twitter.aurora.client.factory.make_client_factory', return_value=mock_api_factory), patch('twitter.aurora.client.api.SchedulerProxy', return_value=mock_scheduler), patch('twitter.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS), patch('twitter.common.app.get_options', return_value=mock_options), patch('twitter.aurora.client.commands.core.get_job_config', return_value=mock_config)) as ( mock_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)