Repository: aurora
Updated Branches:
  refs/heads/master 1a72438f0 -> 59b4d319b


Running task ssh without an instance should pick a random instance

Bugs closed: AURORA-1110

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


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

Branch: refs/heads/master
Commit: 59b4d319b8bb5f48ec3880e36f39527f1498a31c
Parents: 1a72438
Author: JING CHEN <milantr...@gmail.com>
Authored: Fri Sep 30 11:02:26 2016 -0500
Committer: Joshua Cohen <jco...@apache.org>
Committed: Fri Sep 30 11:02:26 2016 -0500

----------------------------------------------------------------------
 .../python/apache/aurora/client/cli/options.py  | 22 +++++++++--
 .../python/apache/aurora/client/cli/task.py     | 18 +++++----
 .../apache/aurora/client/cli/test_task.py       | 41 ++++++++++++++++++++
 3 files changed, 70 insertions(+), 11 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/src/main/python/apache/aurora/client/cli/options.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/cli/options.py 
b/src/main/python/apache/aurora/client/cli/options.py
index 1245ff1..1168703 100644
--- a/src/main/python/apache/aurora/client/cli/options.py
+++ b/src/main/python/apache/aurora/client/cli/options.py
@@ -167,6 +167,18 @@ def parse_options(options):
   return options.split() if options is not None else []
 
 
+def create_instance_argument(help_text):
+  '''Create a CommandOption instance whose type is instance_specifier based on 
help text
+
+  :param help_text: help message for a CommandOption instance
+  :type help_textL string
+  :rtype: CommandOption
+  '''
+  return CommandOption('instance_spec', type=instance_specifier,
+      default=None, metavar="CLUSTER/ROLE/ENV/NAME[/INSTANCES]",
+      help=help_text)
+
+
 BATCH_OPTION = CommandOption('--batch-size', type=int, default=1,
         help='Number of instances to be operate on in one iteration')
 
@@ -212,9 +224,8 @@ INSTANCES_OPTION = CommandOption('--instances', 
type=parse_instances, dest='inst
          'or a range (e.g. 0-2) or any combination of the two (e.g. 
0-2,5,7-9). If not set, '
          'all instances will be acted on.')
 
-INSTANCES_SPEC_ARGUMENT = CommandOption('instance_spec', 
type=instance_specifier,
-    default=None, metavar="CLUSTER/ROLE/ENV/NAME[/INSTANCES]",
-    help=('Fully specified job instance key, in 
CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. '
+INSTANCES_SPEC_ARGUMENT = create_instance_argument(
+    help_text=('Fully specified job instance key, in 
CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. '
         'If INSTANCES is omitted, then all instances will be operated on.'))
 
 
@@ -258,6 +269,10 @@ ROLE_ARGUMENT = CommandOption('role', 
type=parse_qualified_role, metavar='CLUSTE
 ROLE_OPTION = CommandOption('--role', metavar='ROLENAME', default=None,
     help='Name of the user/role')
 
+SSH_INSTANCE_ARGUMENT = create_instance_argument(
+    help_text=('Fully specified job instance key, in 
CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. '
+        'If INSTANCES is omitted, a random instance will picked up for the SSH 
session'))
+
 SSH_USER_OPTION = CommandOption('--ssh-user', '-l', default=None, 
metavar="ssh_username",
     help='ssh as this username instead of the job\'s role')
 
@@ -268,7 +283,6 @@ STRICT_OPTION = CommandOption('--strict', default=False, 
action='store_true',
     help=("Check instances and generate an error for instance ranges in 
parameters "
     "that are larger than the actual set of instances in the job"))
 
-
 TASK_INSTANCE_ARGUMENT = CommandOption('task_instance', 
type=parse_task_instance_key,
     help='A task instance specifier, in the form 
CLUSTER/ROLE/ENV/NAME/INSTANCE')
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/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 a8a4edc..370dd7f 100644
--- a/src/main/python/apache/aurora/client/cli/task.py
+++ b/src/main/python/apache/aurora/client/cli/task.py
@@ -27,11 +27,12 @@ from apache.aurora.client.base import combine_messages
 from apache.aurora.client.cli import EXIT_INVALID_PARAMETER, EXIT_OK, Noun, 
Verb
 from apache.aurora.client.cli.context import AuroraCommandContext
 from apache.aurora.client.cli.options import (
+    ALL_INSTANCES,
     EXECUTOR_SANDBOX_OPTION,
     INSTANCES_SPEC_ARGUMENT,
+    SSH_INSTANCE_ARGUMENT,
     SSH_OPTIONS,
     SSH_USER_OPTION,
-    TASK_INSTANCE_ARGUMENT,
     CommandOption
 )
 from apache.aurora.common.clusters import CLUSTERS
@@ -89,6 +90,7 @@ class SshCommand(Verb):
 
   def get_options(self):
     return [
+        SSH_INSTANCE_ARGUMENT,
         SSH_USER_OPTION,
         SSH_OPTIONS,
         EXECUTOR_SANDBOX_OPTION,
@@ -98,21 +100,23 @@ class SshCommand(Verb):
         CommandOption('--command', '-c', dest='command', type=str, 
default=None,
             metavar="unix_command_line",
             help="Command to execute through the ssh connection."),
-        TASK_INSTANCE_ARGUMENT
     ]
 
   def execute(self, context):
-    (cluster, role, env, name) = context.options.task_instance.jobkey
-    instance = context.options.task_instance.instance
-
+    (cluster, role, env, name) = context.options.instance_spec.jobkey
+    instance = (None if context.options.instance_spec.instance == 
ALL_INSTANCES else
+        set(context.options.instance_spec.instance))
+    if instance is None and context.options.command:
+      raise context.CommandError(EXIT_INVALID_PARAMETER,
+          'INSTANCE must be specified when --command option is given')
     api = context.get_api(cluster)
-    resp = api.query(api.build_query(role, name, env=env, 
instances=set([int(instance)])))
+    resp = api.query(api.build_query(role, name, env=env, instances=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
         len(resp.result.scheduleStatusResult.tasks) == 0):
       raise context.CommandError(EXIT_INVALID_PARAMETER,
-          "Job %s not found" % context.options.task_instance.jobkey)
+          "Job %s not found" % context.options.instance_spec.jobkey)
     first_task = resp.result.scheduleStatusResult.tasks[0]
     remote_cmd = context.options.command or 'bash'
     command = DistributedCommandRunner.substitute(

http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/src/test/python/apache/aurora/client/cli/test_task.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_task.py 
b/src/test/python/apache/aurora/client/cli/test_task.py
index d2233b6..610414f 100644
--- a/src/test/python/apache/aurora/client/cli/test_task.py
+++ b/src/test/python/apache/aurora/client/cli/test_task.py
@@ -152,6 +152,33 @@ class TestSshCommand(AuroraClientCommandTest):
           'cd 
/slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
           'slaverun/sandbox;ls'])
 
+  def test_successful_ssh_no_instance(self):
+    """Test the ssh command when the instance id is not specified."""
+    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+    mock_scheduler_proxy.getTasksStatus.return_value = 
self.create_status_response()
+    sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 
'slaverun'}
+    with contextlib.nested(
+        patch('apache.aurora.client.api.SchedulerProxy', 
return_value=mock_scheduler_proxy),
+        
patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
+            return_value=sandbox_args),
+        patch('subprocess.call', return_value=0)) as (
+            mock_scheduler_proxy_class,
+            mock_runner_args_patch,
+            mock_subprocess):
+      cmd = AuroraCommandLine()
+      cmd.execute(['task', 'ssh', '--ssh-options=-v', 'west/bozo/test/hello'])
+
+      # The status command sends a getTasksStatus query to the scheduler,
+      # and then prints the result.
+      mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
+          jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
+          instanceIds=None,
+          statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, 
ScheduleStatus.RESTARTING,
+              ScheduleStatus.PREEMPTING, ScheduleStatus.DRAINING])))
+      mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost',
+          'cd 
/slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
+          'slaverun/sandbox;bash'])
+
   def test_ssh_job_not_found(self):
     """Test the ssh command when the jobkey parameter specifies a job that 
isn't running."""
     (mock_api, mock_scheduler_proxy) = self.create_mock_api()
@@ -165,3 +192,17 @@ class TestSshCommand(AuroraClientCommandTest):
       result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello/1', 
'--command=ls'])
       assert result == EXIT_INVALID_PARAMETER
       assert mock_subprocess.call_count == 0
+
+  def test_ssh_no_instance_command(self):
+    """Test the ssh command when the jobkey parameter doesn't specify an 
instance."""
+    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+    mock_scheduler_proxy.getTasksStatus.return_value = 
self.create_nojob_status_response()
+    with contextlib.nested(
+        patch('apache.aurora.client.api.SchedulerProxy', 
return_value=mock_scheduler_proxy),
+        patch('subprocess.call', return_value=0)) as (
+            mock_scheduler_proxy_class,
+            mock_subprocess):
+      cmd = AuroraCommandLine()
+      result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello', 
'--command=ls'])
+      assert result == EXIT_INVALID_PARAMETER
+      assert mock_subprocess.call_count == 0

Reply via email to