aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject incubator-aurora git commit: Fixing batched kill task filtering.
Date Sat, 17 Jan 2015 22:34:17 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master 6cf0dea80 -> c37de9ab3


Fixing batched kill task filtering.

Bugs closed: AURORA-996

Reviewed at https://reviews.apache.org/r/29873/


Project: http://git-wip-us.apache.org/repos/asf/incubator-aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-aurora/commit/c37de9ab
Tree: http://git-wip-us.apache.org/repos/asf/incubator-aurora/tree/c37de9ab
Diff: http://git-wip-us.apache.org/repos/asf/incubator-aurora/diff/c37de9ab

Branch: refs/heads/master
Commit: c37de9ab3a32ac287b7b6f0512931f1e310fc465
Parents: 6cf0dea
Author: Maxim Khutornenko <maxim@apache.org>
Authored: Sat Jan 17 14:05:51 2015 -0800
Committer: -l <maxim@apache.org>
Committed: Sat Jan 17 14:05:51 2015 -0800

----------------------------------------------------------------------
 .../python/apache/aurora/client/api/__init__.py |   2 +-
 .../python/apache/aurora/client/cli/context.py  |  22 ++--
 .../python/apache/aurora/client/cli/jobs.py     |  12 +-
 .../python/apache/aurora/client/cli/task.py     |   2 +-
 .../python/apache/aurora/client/cli/update.py   |   2 -
 .../aurora/client/cli/test_command_hooks.py     |  29 ++---
 .../apache/aurora/client/cli/test_create.py     |  35 ++----
 .../apache/aurora/client/cli/test_kill.py       |  66 +++++++++--
 .../apache/aurora/client/cli/test_plugins.py    |  21 +---
 .../apache/aurora/client/cli/test_restart.py    |  17 +++
 .../apache/aurora/client/cli/test_supdate.py    |  14 +++
 .../apache/aurora/client/cli/test_update.py     |  19 ++-
 .../python/apache/aurora/client/cli/util.py     | 118 ++++++++++---------
 13 files changed, 209 insertions(+), 150 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/main/python/apache/aurora/client/api/__init__.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/__init__.py b/src/main/python/apache/aurora/client/api/__init__.py
index 64f0804..07d1b1c 100644
--- a/src/main/python/apache/aurora/client/api/__init__.py
+++ b/src/main/python/apache/aurora/client/api/__init__.py
@@ -121,7 +121,7 @@ class AuroraClientAPI(object):
     return self.query_no_configs(job_key.to_thrift_query())
 
   @classmethod
-  def build_query(cls, role, job, instances=None, statuses=LIVE_STATES, env=None):
+  def build_query(cls, role, job, env=None, instances=None, statuses=LIVE_STATES):
     return TaskQuery(jobKeys=[JobKey(role=role, environment=env, name=job)],
                      statuses=statuses,
                      instanceIds=instances)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/main/python/apache/aurora/client/cli/context.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/cli/context.py b/src/main/python/apache/aurora/client/cli/context.py
index 93587c6..45fb404 100644
--- a/src/main/python/apache/aurora/client/cli/context.py
+++ b/src/main/python/apache/aurora/client/cli/context.py
@@ -34,6 +34,7 @@ from apache.aurora.client.factory import make_client
 from apache.aurora.common.aurora_job_key import AuroraJobKey
 from apache.aurora.common.clusters import CLUSTERS
 
+from gen.apache.aurora.api.constants import ACTIVE_STATES
 from gen.apache.aurora.api.ttypes import ResponseCode
 
 # Utility type, representing job keys with wildcards.
@@ -202,24 +203,27 @@ class AuroraCommandContext(Context):
       return jobs
 
   def get_job_status(self, key):
-    """Returns a list of task instances running under the job."""
+    """Returns a list of task instances."""
     api = self.get_api(key.cluster)
     resp = api.check_status(key)
     self.log_response_and_raise(resp, err_code=EXIT_INVALID_PARAMETER)
-    return resp.result.scheduleStatusResult.tasks or None
+    return resp.result.scheduleStatusResult.tasks
 
   def get_active_instances(self, key):
     """Returns a list of the currently active instances of a job"""
-    return [task.assignedTask.instanceId for task in self.get_job_status(key)]
+    api = self.get_api(key.cluster)
+    resp = api.query_no_configs(
+        api.build_query(key.role, key.name, env=key.env, statuses=ACTIVE_STATES))
+    self.log_response_and_raise(resp, err_code=EXIT_INVALID_PARAMETER)
+    return resp.result.scheduleStatusResult.tasks
 
   def verify_instances_option_validity(self, jobkey, instances):
-    """Given a jobkey, does a getTasksStatus, and then checks that the specified instances
-    are valid for the job.
-    """
-    active_instances = self.get_active_instances(jobkey)
-    if max(active_instances) < max(instances):
+    """Verifies all provided job instances are currently active."""
+    active = set(task.assignedTask.instanceId for task in self.get_active_instances(jobkey)
or [])
+    unrecognized = set(instances) - active
+    if unrecognized:
       raise self.CommandError(EXIT_INVALID_PARAMETER,
-          "Invalid shards parameter: %s only has %s shards" % (jobkey, max(active_instances)))
+          "Invalid instance parameter: %s" % (list(unrecognized)))
 
   def timestamp_to_string(self, timestamp):
     return time.ctime(timestamp)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/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 508c9be..1ac7145 100644
--- a/src/main/python/apache/aurora/client/cli/jobs.py
+++ b/src/main/python/apache/aurora/client/cli/jobs.py
@@ -198,7 +198,7 @@ class DiffCommand(Verb):
       env = config.environment()
       name = config.name()
     api = context.get_api(cluster)
-    resp = api.query(api.build_query(role, name, statuses=ACTIVE_STATES, env=env))
+    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:
@@ -318,7 +318,7 @@ class AbstractKillCommand(Verb):
   def kill_in_batches(self, context, job, instances_arg):
     api = context.get_api(job.cluster)
     # query the job, to get the list of active instances.
-    tasks = context.get_job_status(job)
+    tasks = context.get_active_instances(job)
     if tasks is None or len(tasks) == 0:
       context.print_err("No tasks to kill found for job %s" % job)
       return EXIT_INVALID_PARAMETER
@@ -611,7 +611,7 @@ class StatusCommand(Verb):
     result = []
     for jk in jobkeys:
       job_tasks = context.get_job_status(jk)
-      if job_tasks is None or job_tasks is []:
+      if not job_tasks:
         logging.info("No tasks were found for jobkey %s" % jk)
         continue
       active_tasks = sorted([t for t in job_tasks if is_active(t)],
@@ -646,7 +646,7 @@ class StatusCommand(Verb):
 
   def execute(self, context):
     jobs = context.get_jobs_matching_key(context.options.jobspec)
-    if jobs is None or jobs == []:
+    if not jobs:
       return self._print_jobs_not_found()
 
     result = self.get_status_for_jobs(jobs, context)
@@ -688,7 +688,7 @@ class UpdateCommand(Verb):
     # Get the current job status, so that we can check if there's anything
     # dangerous about this update.
     resp = api.query_no_configs(api.build_query(config.role(), config.name(),
-        statuses=ACTIVE_STATES, env=config.environment()))
+        env=config.environment(), statuses=ACTIVE_STATES))
     context.log_response_and_raise(resp, err_msg="Server could not find running job to update")
     remote_tasks = [t.assignedTask.task for t in resp.result.scheduleStatusResult.tasks]
     # for determining if an update is dangerous, we estimate the scope of the change
@@ -715,8 +715,6 @@ class UpdateCommand(Verb):
     job = context.options.instance_spec.jobkey
     instances = (None if context.options.instance_spec.instance == ALL_INSTANCES else
         context.options.instance_spec.instance)
-    if instances is not None and context.options.strict:
-      context.verify_instances_option_validity(job, instances)
     config = context.get_job_config(job, context.options.config_file)
     api = context.get_api(config.cluster())
     if not context.options.force:

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/main/python/apache/aurora/client/cli/task.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/cli/task.py b/src/main/python/apache/aurora/client/cli/task.py
index 2154898..b541d03 100644
--- a/src/main/python/apache/aurora/client/cli/task.py
+++ b/src/main/python/apache/aurora/client/cli/task.py
@@ -97,7 +97,7 @@ class SshCommand(Verb):
     instance = context.options.task_instance.instance
 
     api = context.get_api(cluster)
-    resp = api.query(api.build_query(role, name, set([int(instance)]), env=env))
+    resp = api.query(api.build_query(role, name, env=env, instances=set([int(instance)])))
     context.log_response_and_raise(resp,
         err_msg=('Unable to get information about instance: %s' % combine_messages(resp)))
     if (resp.result.scheduleStatusResult.tasks is None or

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/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 a161732..56f9228 100644
--- a/src/main/python/apache/aurora/client/cli/update.py
+++ b/src/main/python/apache/aurora/client/cli/update.py
@@ -72,8 +72,6 @@ class StartUpdate(Verb):
     job = context.options.instance_spec.jobkey
     instances = (None if context.options.instance_spec.instance == ALL_INSTANCES else
         context.options.instance_spec.instance)
-    if instances is not None and context.options.strict:
-      context.verify_instances_option_validity(job, instances)
     config = context.get_job_config(job, context.options.config_file)
     if config.raw().has_cron_schedule():
       raise context.CommandError(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_command_hooks.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_command_hooks.py b/src/test/python/apache/aurora/client/cli/test_command_hooks.py
index e8432a1..2130f1f 100644
--- a/src/test/python/apache/aurora/client/cli/test_command_hooks.py
+++ b/src/test/python/apache/aurora/client/cli/test_command_hooks.py
@@ -22,13 +22,10 @@ from apache.aurora.config import AuroraConfig
 from .util import AuroraClientCommandTest, FakeAuroraCommandContext
 
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
     JobKey,
     Result,
-    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
-    TaskEvent,
     TaskQuery
 )
 
@@ -66,30 +63,20 @@ class HookForTesting(CommandHook):
 class TestClientCreateCommand(AuroraClientCommandTest):
 
   @classmethod
-  def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    return ScheduledTask(
-        status=status,
-        assignedTask=AssignedTask(
-          taskId=task_id,
-          instanceId=instance_id),
-        taskEvents=[TaskEvent(timestamp=initial_time)]
-    )
-
-  @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
-    mock_query_result = cls.create_simple_success_response()
+    query_result = cls.create_simple_success_response()
     if scheduleStatus == ScheduleStatus.INIT:
       # status query result for before job is launched.
       tasks = []
     else:
-      mock_task_one = cls.create_mock_task("hello", 0, 1000, scheduleStatus)
-      mock_task_two = cls.create_mock_task("hello", 1, 1004, scheduleStatus)
-      tasks = [mock_task_one, mock_task_two]
-    mock_query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
-    return mock_query_result
+      task_one = cls.create_scheduled_task(0, initial_time=1000, status=scheduleStatus)
+      task_two = cls.create_scheduled_task(1, initial_time=1004, status=scheduleStatus)
+      tasks = [task_one, task_two]
+    query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
+    return query_result
 
   @classmethod
-  def create_mock_query(cls):
+  def create_query(cls):
     return TaskQuery(
         jobKeys=[JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=cls.TEST_JOB)])
 
@@ -115,7 +102,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
     GlobalCommandHookRegistry.register_command_hook(command_hook)
     mock_context = FakeAuroraCommandContext()
     with patch("apache.aurora.client.cli.jobs.Job.create_context", return_value=mock_context):
-      mock_query = self.create_mock_query()
+      mock_query = self.create_query()
       mock_context.add_expected_status_query_result(
           self.create_mock_status_query_result(ScheduleStatus.INIT))
       mock_context.add_expected_status_query_result(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_create.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_create.py b/src/test/python/apache/aurora/client/cli/test_create.py
index 18b5838..a65aab7 100644
--- a/src/test/python/apache/aurora/client/cli/test_create.py
+++ b/src/test/python/apache/aurora/client/cli/test_create.py
@@ -39,14 +39,11 @@ from .util import (
 )
 
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
     JobKey,
     ResponseCode,
     Result,
-    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
-    TaskEvent,
     TaskQuery
 )
 
@@ -88,28 +85,16 @@ class TestCreateJobCommand(AuroraClientCommandTest):
 class TestClientCreateCommand(AuroraClientCommandTest):
 
   @classmethod
-  def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    mock_task = create_autospec(spec=ScheduledTask, instance=True)
-    mock_task.assignedTask = create_autospec(spec=AssignedTask, instance=True)
-    mock_task.assignedTask.taskId = task_id
-    mock_task.assignedTask.instanceId = instance_id
-    mock_task.status = status
-    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
-    mock_task_event.timestamp = initial_time
-    mock_task.taskEvents = [mock_task_event]
-    return mock_task
-
-  @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
-    mock_query_result = cls.create_simple_success_response()
-    mock_query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[
-        cls.create_mock_task('hello', 0, 1000, scheduleStatus),
-        cls.create_mock_task('hello', 1, 1004, scheduleStatus)
+    query_result = cls.create_simple_success_response()
+    query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[
+        cls.create_scheduled_task(0, initial_time=1000, status=scheduleStatus),
+        cls.create_scheduled_task(1, initial_time=1004, status=scheduleStatus)
     ]))
-    return mock_query_result
+    return query_result
 
   @classmethod
-  def create_mock_query(cls):
+  def create_query(cls):
     return TaskQuery(
         jobKeys=[JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=cls.TEST_JOB)])
 
@@ -149,7 +134,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
       # After making the client, create sets up a job monitor.
       # The monitor uses TaskQuery to get the tasks. It's called at least twice:once before
       # the job is created, and once after. So we need to set up mocks for the query results.
-      mock_query = self.create_mock_query()
+      mock_query = self.create_query()
       mock_context.add_expected_status_query_result(
         self.create_mock_status_query_result(ScheduleStatus.PENDING))
       mock_context.add_expected_status_query_result(
@@ -177,7 +162,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
         #              combinations did not produce the desired effect. Investigate why (AURORA-510)
         patch('threading._Event.wait'),
         patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
-      mock_query = self.create_mock_query()
+      mock_query = self.create_query()
       mock_context.add_expected_status_query_result(
         self.create_mock_status_query_result(ScheduleStatus.PENDING))
       mock_context.add_expected_status_query_result(
@@ -205,7 +190,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
     with contextlib.nested(
         patch('threading._Event.wait'),
         patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
-      mock_query = self.create_mock_query()
+      mock_query = self.create_query()
       for result in [ScheduleStatus.PENDING, ScheduleStatus.PENDING, ScheduleStatus.RUNNING]:
         mock_context.add_expected_status_query_result(self.create_mock_status_query_result(result))
       api = mock_context.get_api('west')
@@ -381,7 +366,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
     with contextlib.nested(
         patch('threading._Event.wait'),
         patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
-      mock_query = self.create_mock_query()
+      mock_query = self.create_query()
       mock_context.add_expected_status_query_result(
         self.create_mock_status_query_result(ScheduleStatus.PENDING))
       mock_context.add_expected_status_query_result(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_kill.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_kill.py b/src/test/python/apache/aurora/client/cli/test_kill.py
index b475d73..7aad34a 100644
--- a/src/test/python/apache/aurora/client/cli/test_kill.py
+++ b/src/test/python/apache/aurora/client/cli/test_kill.py
@@ -33,6 +33,7 @@ from .util import (
     mock_verb_options
 )
 
+from gen.apache.aurora.api.constants import ACTIVE_STATES
 from gen.apache.aurora.api.ttypes import (
     JobKey,
     ResponseCode,
@@ -92,9 +93,9 @@ class TestKillCommand(AuroraClientCommandTest):
     fake_context = FakeAuroraCommandContext()
     fake_context.set_options(mock_options)
 
-    fake_context.add_expected_status_query_result(
-      AuroraClientCommandTest.create_status_call_result(
-        AuroraClientCommandTest.create_mock_task(1, ScheduleStatus.KILLED)))
+    fake_context.add_expected_query_result(
+      AuroraClientCommandTest.create_query_call_result(
+        AuroraClientCommandTest.create_scheduled_task(1, ScheduleStatus.RUNNING)))
 
     mock_api = fake_context.get_api('test')
     mock_api.kill_job.return_value = AuroraClientCommandTest.create_blank_response(
@@ -106,6 +107,44 @@ class TestKillCommand(AuroraClientCommandTest):
     mock_api.kill_job.assert_called_once_with(jobkey, mock_options.instance_spec.instance)
     self.assert_lock_message(fake_context)
 
+  def test_kill_inactive_instance_spec(self):
+    """Verify the instance spec is validated in a batched kill."""
+    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
+    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.value.message == "Invalid instance parameter: [1]"
+
+  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"
+
 
 class TestClientKillCommand(AuroraClientCommandTest):
   @classmethod
@@ -145,8 +184,8 @@ class TestClientKillCommand(AuroraClientCommandTest):
 
   @classmethod
   def assert_query(cls, fake_api):
-    calls = [call(cls.TEST_JOBKEY)]
-    assert fake_api.check_status.mock_calls == calls
+    calls = [call(TaskQuery(jobKeys=[cls.TEST_JOBKEY.to_thrift()], statuses=ACTIVE_STATES))]
+    assert fake_api.query_no_configs.mock_calls == calls
 
   def test_killall_job(self):
     """Test killall client-side API logic."""
@@ -181,7 +220,8 @@ class TestClientKillCommand(AuroraClientCommandTest):
 
       api = mock_context.get_api('west')
       api.kill_job.return_value = self.create_simple_success_response()
-      mock_context.add_expected_status_query_result(self.create_status_call_result())
+      mock_context.add_expected_query_result(
+          self.create_query_call_result(), job_key=self.TEST_JOBKEY)
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
@@ -264,7 +304,9 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)) as (_, m,
_):
       api = mock_context.get_api('west')
-      mock_context.add_expected_status_query_result(self.create_status_call_result())
+      mock_context.add_expected_query_result(
+          self.create_query_call_result(), job_key=self.TEST_JOBKEY)
+
       api.kill_job.return_value = self.create_simple_success_response()
 
       with temporary_file() as fp:
@@ -286,7 +328,9 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)):
       api = mock_context.get_api('west')
-      mock_context.add_expected_status_query_result(self.create_status_call_result())
+      mock_context.add_expected_query_result(
+          self.create_query_call_result(), job_key=self.TEST_JOBKEY)
+
       api.kill_job.return_value = self.create_simple_success_response()
 
       with temporary_file() as fp:
@@ -310,7 +354,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
       # set up an empty instance list in the getTasksWithoutConfigs response
       status_response = self.create_simple_success_response()
       status_response.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[]))
-      mock_context.add_expected_status_query_result(status_response)
+      mock_context.add_expected_query_result(status_response)
       api.kill_job.return_value = self.create_simple_success_response()
 
       with temporary_file() as fp:
@@ -348,7 +392,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=self.get_monitor_mock()),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)):
       api = mock_context.get_api('west')
-      mock_context.add_expected_status_query_result(self.create_status_call_result())
+      mock_context.add_expected_query_result(self.create_query_call_result())
       api.kill_job.return_value = self.create_simple_success_response()
 
       with temporary_file() as fp:
@@ -371,7 +415,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)):
       api = mock_context.get_api('west')
-      mock_context.add_expected_status_query_result(self.create_status_call_result())
+      mock_context.add_expected_query_result(self.create_query_call_result())
       api.kill_job.return_value = self.create_simple_success_response()
 
       with temporary_file() as fp:

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_plugins.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_plugins.py b/src/test/python/apache/aurora/client/cli/test_plugins.py
index aa45851..a545fec 100644
--- a/src/test/python/apache/aurora/client/cli/test_plugins.py
+++ b/src/test/python/apache/aurora/client/cli/test_plugins.py
@@ -12,7 +12,7 @@
 # limitations under the License.
 #
 
-from mock import create_autospec, patch
+from mock import patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.cli import ConfigurationPlugin
@@ -23,13 +23,10 @@ from apache.aurora.config import AuroraConfig
 from .util import AuroraClientCommandTest, FakeAuroraCommandContext
 
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
     JobKey,
     Result,
-    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
-    TaskEvent,
     TaskQuery
 )
 
@@ -61,26 +58,14 @@ class BogusPlugin(ConfigurationPlugin):
 class TestPlugins(AuroraClientCommandTest):
 
   @classmethod
-  def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    mock_task = create_autospec(spec=ScheduledTask, instance=True)
-    mock_task.assignedTask = create_autospec(spec=AssignedTask, instance=True)
-    mock_task.assignedTask.taskId = task_id
-    mock_task.assignedTask.instanceId = instance_id
-    mock_task.status = status
-    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
-    mock_task_event.timestamp = initial_time
-    mock_task.taskEvents = [mock_task_event]
-    return mock_task
-
-  @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
     mock_query_result = cls.create_simple_success_response()
     if scheduleStatus == ScheduleStatus.INIT:
       # status query result for before job is launched.
       tasks = []
     else:
-      mock_task_one = cls.create_mock_task('hello', 0, 1000, scheduleStatus)
-      mock_task_two = cls.create_mock_task('hello', 1, 1004, scheduleStatus)
+      mock_task_one = cls.create_scheduled_task(0, initial_time=1000, status=scheduleStatus)
+      mock_task_two = cls.create_scheduled_task(1, initial_time=1004, status=scheduleStatus)
       tasks = [mock_task_one, mock_task_two]
     mock_query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
     return mock_query_result

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_restart.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_restart.py b/src/test/python/apache/aurora/client/cli/test_restart.py
index a532ead..aa389b4 100644
--- a/src/test/python/apache/aurora/client/cli/test_restart.py
+++ b/src/test/python/apache/aurora/client/cli/test_restart.py
@@ -62,6 +62,23 @@ class TestRestartJobCommand(AuroraClientCommandTest):
       updater_config, mock_options.healthcheck_interval_seconds, config=None)
     self.assert_lock_message(fake_context)
 
+  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]"
+
 
 class TestRestartCommand(AuroraClientCommandTest):
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/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 9378b49..5c3e22f 100644
--- a/src/test/python/apache/aurora/client/cli/test_supdate.py
+++ b/src/test/python/apache/aurora/client/cli/test_supdate.py
@@ -88,6 +88,20 @@ class TestStartUpdateCommand(AuroraClientCommandTest):
     with pytest.raises(Context.CommandError):
       self._command.execute(self._fake_context)
 
+  def test_update_no_active_instance_check(self):
+    self._mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
+    self._mock_options.strict = True
+
+    mock_config = self.create_mock_config()
+    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)
+
+    self._mock_api.start_job_update.assert_called_once_with(
+      mock_config,
+      self._mock_options.instance_spec.instance)
+
 
 class TestUpdateCommand(AuroraClientCommandTest):
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/test_update.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_update.py b/src/test/python/apache/aurora/client/cli/test_update.py
index c12b32e..7985436 100644
--- a/src/test/python/apache/aurora/client/cli/test_update.py
+++ b/src/test/python/apache/aurora/client/cli/test_update.py
@@ -70,7 +70,7 @@ class TestJobUpdateCommand(AuroraClientCommandTest):
   def test_update_with_lock(self):
     mock_config = self.create_mock_config()
     self._fake_context.get_job_config = Mock(return_value=mock_config)
-    self._mock_api.update_job.return_value = AuroraClientCommandTest.create_blank_response(
+    self._mock_api.update_job.return_value = self.create_blank_response(
         ResponseCode.LOCK_ERROR, "Error.")
 
     with pytest.raises(Context.CommandError):
@@ -86,7 +86,7 @@ class TestJobUpdateCommand(AuroraClientCommandTest):
     mock_config = self.create_mock_config()
     self._fake_context.get_job_config = Mock(return_value=mock_config)
     error = "Error printed once."
-    self._mock_api.update_job.return_value = AuroraClientCommandTest.create_blank_response(
+    self._mock_api.update_job.return_value = self.create_blank_response(
         ResponseCode.INVALID_REQUEST,
         error)
 
@@ -99,6 +99,21 @@ class TestJobUpdateCommand(AuroraClientCommandTest):
       self._mock_options.instance_spec.instance)
     assert self._fake_context.get_err() == ["Update failed due to error:", "\t%s" % error]
 
+  def test_update_no_active_instance_check(self):
+    self._mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
+    self._mock_options.strict = True
+
+    mock_config = self.create_mock_config()
+    self._fake_context.get_job_config = Mock(return_value=mock_config)
+    self._mock_api.update_job.return_value = self.create_simple_success_response()
+
+    self._command.execute(self._fake_context)
+
+    self._mock_api.update_job.assert_called_once_with(
+      mock_config,
+      self._mock_options.healthcheck_interval_seconds,
+      self._mock_options.instance_spec.instance)
+
 
 class TestUpdateCommand(AuroraClientCommandTest):
   class FakeSchedulerMux(SchedulerMux):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/c37de9ab/src/test/python/apache/aurora/client/cli/util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/util.py b/src/test/python/apache/aurora/client/cli/util.py
index 147d418..5b6207d 100644
--- a/src/test/python/apache/aurora/client/cli/util.py
+++ b/src/test/python/apache/aurora/client/cli/util.py
@@ -26,6 +26,7 @@ from apache.aurora.common.clusters import Clusters
 
 from ...api_util import SchedulerProxyApiSpec, SchedulerThriftApiSpec
 
+from gen.apache.aurora.api.constants import ACTIVE_STATES
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
     ExecutorConfig,
@@ -39,7 +40,8 @@ from gen.apache.aurora.api.ttypes import (
     ScheduleStatus,
     ScheduleStatusResult,
     TaskConfig,
-    TaskEvent
+    TaskEvent,
+    TaskQuery
 )
 
 
@@ -74,7 +76,7 @@ class FakeAuroraCommandContext(AuroraCommandContext):
     super(FakeAuroraCommandContext, self).__init__()
     self.status = []
     self.fake_api = self.create_mock_api()
-    self.task_status = []
+    self.task_result = []
     self.showed_urls = []
     self.out = []
     self.err = []
@@ -121,10 +123,20 @@ class FakeAuroraCommandContext(AuroraCommandContext):
     return "YYYY-MM-DD HH:MM:SS"
 
   def add_expected_status_query_result(self, expected_result):
-    self.task_status.append(expected_result)
+    self.add_task_result(expected_result)
+    self.fake_api.check_status.side_effect = self.task_result
+
+  def add_expected_query_result(self, expected_result, job_key=None):
+    self.add_task_result(expected_result)
+    self.fake_api.query_no_configs.side_effect = self.task_result
+    if job_key:
+      self.fake_api.build_query.return_value = TaskQuery(
+          jobKeys=[job_key.to_thrift()], statuses=ACTIVE_STATES)
+
+  def add_task_result(self, expected_result):
+    self.task_result.append(expected_result)
     # each call adds an expected query result, in order.
-    self.fake_api.scheduler_proxy.getTasksWithoutConfigs.side_effect = self.task_status
-    self.fake_api.check_status.side_effect = self.task_status
+    self.fake_api.scheduler_proxy.getTasksWithoutConfigs.side_effect = self.task_result
 
 
 class AuroraClientCommandTest(unittest.TestCase):
@@ -162,69 +174,69 @@ class AuroraClientCommandTest(unittest.TestCase):
     return mock_api_factory, mock_scheduler_client
 
   @classmethod
-  def create_status_call_result(cls, mock_task=None):
-    status_response = cls.create_simple_success_response()
-    schedule_status = create_autospec(spec=ScheduleStatusResult, instance=True)
-    status_response.result = Result(scheduleStatusResult=schedule_status)
-    # This should be a list of ScheduledTask's.
-    schedule_status.tasks = []
-    if mock_task is None:
+  def create_query_call_result(cls, task=None):
+    status_response = cls.create_empty_task_result()
+    if task is None:
       for i in range(20):
-        schedule_status.tasks.append(cls.create_mock_task(i))
+        status_response.result.scheduleStatusResult.tasks.append(cls.create_scheduled_task(i))
     else:
-      schedule_status.tasks.append(mock_task)
+      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_mock_task(cls, instance_id, status=ScheduleStatus.RUNNING):
-    mock_task = create_autospec(spec=ScheduledTask, instance=True)
-    mock_task.assignedTask = create_autospec(spec=AssignedTask, instance=True)
-    mock_task.assignedTask.instanceId = instance_id
-    mock_task.assignedTask.taskId = "Task%s" % instance_id
-    mock_task.assignedTask.slaveId = "Slave%s" % instance_id
-    mock_task.assignedTask.task = create_autospec(spec=TaskConfig, instance=True)
-    mock_task.slaveHost = "Slave%s" % instance_id
-    mock_task.status = status
-    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
-    mock_task_event.timestamp = 1000
-    mock_task.taskEvents = [mock_task_event]
-    return mock_task
+  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()
-      task.failure_count = 0
-      task.assignedTask = AssignedTask()
-      task.assignedTask.taskId = 1287391823
-      task.assignedTask.slaveHost = 'slavehost'
-      task.assignedTask.task = TaskConfig()
-      task.assignedTask.task.maxTaskFailures = 1
-      task.assignedTask.task.executorConfig = ExecutorConfig()
-      task.assignedTask.task.executorConfig.data = 'fake data'
-      task.assignedTask.task.metadata = []
-      task.assignedTask.task.job = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      task.assignedTask.task.owner = Identity(role=cls.TEST_ROLE)
-      task.assignedTask.task.environment = cls.TEST_ENV
-      task.assignedTask.task.jobName = name
-      task.assignedTask.task.numCpus = 2
-      task.assignedTask.task.ramMb = 2
-      task.assignedTask.task.diskMb = 2
-      task.assignedTask.instanceId = 4237894
-      task.assignedTask.assignedPorts = {}
-      task.status = ScheduleStatus.RUNNING
-      event = TaskEvent()
-      event.timestamp = 28234726395
-      event.status = ScheduleStatus.RUNNING
-      event.message = "Hi there"
-      task.taskEvents = [event]
+      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_status_call_result()
+    status_response = cls.create_query_call_result()
     scheduler.getTasksWithoutConfigs.return_value = status_response
 
   @classmethod


Mime
View raw message