You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@aurora.apache.org by ma...@apache.org on 2014/10/24 01:07:49 UTC

[2/2] git commit: Preparing for Identity struct deprecation (client and executor).

Preparing for Identity struct deprecation (client and executor).

Bugs closed: AURORA-84

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


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

Branch: refs/heads/master
Commit: 06935c04200fa78aea2ed60ee49ed234d796c6bd
Parents: 18ae0ab
Author: Maxim Khutornenko <ma...@apache.org>
Authored: Thu Oct 23 16:00:38 2014 -0700
Committer: Maxim Khutornenko <ma...@apache.org>
Committed: Thu Oct 23 16:00:38 2014 -0700

----------------------------------------------------------------------
 .../aurora/client/api/instance_watcher.py       |  6 +-
 src/main/python/apache/aurora/client/api/sla.py | 19 ++---
 .../python/apache/aurora/client/cli/jobs.py     | 10 +--
 .../python/apache/aurora/client/cli/task.py     | 11 +--
 .../apache/aurora/client/commands/admin.py      |  4 +-
 .../apache/aurora/client/commands/core.py       | 20 ++---
 .../python/apache/aurora/client/commands/ssh.py | 11 +--
 src/main/python/apache/aurora/config/thrift.py  |  1 +
 .../apache/aurora/executor/aurora_executor.py   |  2 +-
 .../apache/aurora/executor/common/announcer.py  |  7 +-
 .../aurora/executor/thermos_task_runner.py      |  2 +-
 .../aurora/client/api/test_instance_watcher.py  | 18 ++---
 .../python/apache/aurora/client/api/test_sla.py | 12 ++-
 .../aurora/client/cli/test_api_from_cli.py      | 37 +--------
 .../apache/aurora/client/cli/test_diff.py       | 43 +---------
 .../apache/aurora/client/cli/test_status.py     | 84 +++++++-------------
 .../apache/aurora/client/cli/test_task_run.py   | 69 +---------------
 .../apache/aurora/client/cli/test_update.py     | 13 ++-
 .../python/apache/aurora/client/cli/util.py     | 36 ++++++++-
 .../apache/aurora/client/commands/test_diff.py  |  9 ++-
 .../apache/aurora/client/commands/test_ssh.py   | 10 ++-
 .../aurora/client/commands/test_status.py       | 10 ++-
 .../aurora/client/commands/test_update.py       | 14 +++-
 .../python/apache/aurora/config/test_thrift.py  | 10 ++-
 .../aurora/executor/common/test_announcer.py    | 12 ++-
 .../aurora/executor/test_thermos_executor.py    |  3 +-
 26 files changed, 187 insertions(+), 286 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/client/api/instance_watcher.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/instance_watcher.py b/src/main/python/apache/aurora/client/api/instance_watcher.py
index b390aa8..6ed8154 100644
--- a/src/main/python/apache/aurora/client/api/instance_watcher.py
+++ b/src/main/python/apache/aurora/client/api/instance_watcher.py
@@ -20,7 +20,7 @@ from twitter.common import log
 from .health_check import StatusHealthCheck
 from .task_util import StatusMuxHelper
 
-from gen.apache.aurora.api.ttypes import Identity, ScheduleStatus, TaskQuery
+from gen.apache.aurora.api.ttypes import ScheduleStatus, TaskQuery
 
 
 class Instance(object):
@@ -133,9 +133,7 @@ class InstanceWatcher(object):
 
   def _create_query(self, instance_ids):
     query = TaskQuery()
-    query.owner = Identity(role=self._job_key.role)
-    query.environment = self._job_key.environment
-    query.jobName = self._job_key.name
+    query.jobKeys = set([self._job_key])
     query.statuses = set([ScheduleStatus.RUNNING])
     query.instanceIds = instance_ids
     return query

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/client/api/sla.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/sla.py b/src/main/python/apache/aurora/client/api/sla.py
index b9b6468..5855685 100644
--- a/src/main/python/apache/aurora/client/api/sla.py
+++ b/src/main/python/apache/aurora/client/api/sla.py
@@ -22,7 +22,7 @@ from apache.aurora.client.base import DEFAULT_GROUPING, format_response, group_h
 from apache.aurora.common.aurora_job_key import AuroraJobKey
 
 from gen.apache.aurora.api.constants import LIVE_STATES
-from gen.apache.aurora.api.ttypes import Identity, ResponseCode, ScheduleStatus, TaskQuery
+from gen.apache.aurora.api.ttypes import ResponseCode, ScheduleStatus, TaskQuery
 
 
 def job_key_from_scheduled(task, cluster):
@@ -32,28 +32,25 @@ def job_key_from_scheduled(task, cluster):
   task -- ScheduledTask to get job key from.
   cluster -- Cluster the task belongs to.
   """
+  config = task.assignedTask.task
   return AuroraJobKey(
       cluster=cluster.name,
-      role=task.assignedTask.task.owner.role,
-      env=task.assignedTask.task.environment,
-      name=task.assignedTask.task.jobName
+      role=config.job.role if config.job else config.owner.role,
+      env=config.job.environment if config.job else config.environment,
+      name=config.job.name if config.job else config.jobName
   )
 
 
-def task_query(job_key=None, hosts=None, job_keys=None):
+def task_query(hosts=None, job_keys=None):
   """Creates TaskQuery optionally scoped by a job(s) or hosts.
 
   Arguments:
-  job_key -- AuroraJobKey to scope the query by.
   hosts -- list of hostnames to scope the query by.
   job_keys -- list of AuroraJobKeys to scope the query by.
   """
   return TaskQuery(
-      owner=Identity(role=job_key.role) if job_key else None,
-      environment=job_key.env if job_key else None,
-      jobName=job_key.name if job_key else None,
       slaveHosts=set(hosts) if hosts else None,
-      jobKeys=set(k.to_thrift() for k in job_keys) if job_keys else None,
+      jobKeys=[k.to_thrift() for k in job_keys] if job_keys else None,
       statuses=LIVE_STATES)
 
 
@@ -302,7 +299,7 @@ class Sla(object):
     Arguments:
     job_key -- job to create a task uptime vector for.
     """
-    return JobUpTimeSlaVector(self._get_tasks(task_query(job_key=job_key)))
+    return JobUpTimeSlaVector(self._get_tasks(task_query(job_keys=[job_key])))
 
   def get_domain_uptime_vector(self, cluster, min_instance_count, hosts=None):
     """Returns a DomainUpTimeSlaVector object with all available job uptimes.

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/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 493ba21..625cb80 100644
--- a/src/main/python/apache/aurora/client/cli/jobs.py
+++ b/src/main/python/apache/aurora/client/cli/jobs.py
@@ -562,12 +562,12 @@ class StatusCommand(Verb):
       task_info = assigned_task.task
       task_strings = []
       task_strings.append("\tTask role: %s, env: %s, name: %s, instance: %s, status: %s on %s" %
-             (scheduled_task.assignedTask.task.owner.role,
-              scheduled_task.assignedTask.task.environment,
-              scheduled_task.assignedTask.task.jobName,
-              scheduled_task.assignedTask.instanceId,
+             (task_info.job.role if task_info.job else task_info.owner.role,
+              task_info.job.environment if task_info.job else task_info.environment,
+              task_info.job.name if task_info.job else task_info.jobName,
+              assigned_task.instanceId,
               ScheduleStatus._VALUES_TO_NAMES[scheduled_task.status],
-              scheduled_task.assignedTask.slaveHost))
+              assigned_task.slaveHost))
 
       if task_info:
         task_strings.append("""\t  cpus: %s, ram: %s MB, disk: %s MB""" % (

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/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 c41484b..37e0141 100644
--- a/src/main/python/apache/aurora/client/cli/task.py
+++ b/src/main/python/apache/aurora/client/cli/task.py
@@ -114,8 +114,9 @@ class SshCommand(Verb):
         api.cluster, executor_sandbox=context.options.executor_sandbox)
 
     ssh_command = ['ssh', '-t']
-    role = first_task.assignedTask.task.owner.role
-    slave_host = first_task.assignedTask.slaveHost
+    assigned = first_task.assignedTask
+    role = assigned.task.job.role if assigned.task.job else assigned.task.owner.role
+    slave_host = assigned.slaveHost
 
     for tunnel in context.options.tunnels:
       try:
@@ -124,11 +125,11 @@ class SshCommand(Verb):
       except ValueError:
         raise context.CommandError(EXIT_INVALID_PARAMETER,
             'Could not parse tunnel: %s.  Must be of form PORT:NAME' % tunnel)
-      if name not in first_task.assignedTask.assignedPorts:
+      if name not in assigned.assignedPorts:
         raise context.CommandError(EXIT_INVALID_PARAMETER,
-            'Task %s has no port named %s' % (first_task.assignedTask.taskId, name))
+            'Task %s has no port named %s' % (assigned.taskId, name))
       ssh_command += [
-          '-L', '%d:%s:%d' % (port, slave_host, first_task.assignedTask.assignedPorts[name])]
+          '-L', '%d:%s:%d' % (port, slave_host, assigned.assignedPorts[name])]
 
     ssh_command += ['%s@%s' % (context.options.ssh_user or role, slave_host), command]
     return subprocess.call(ssh_command)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/client/commands/admin.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/admin.py b/src/main/python/apache/aurora/client/commands/admin.py
index deee025..9719a58 100644
--- a/src/main/python/apache/aurora/client/commands/admin.py
+++ b/src/main/python/apache/aurora/client/commands/admin.py
@@ -268,8 +268,8 @@ def scheduler_print_recovery_tasks(cluster):
   for task in resp.result.queryRecoveryResult.tasks:
     assigned = task.assignedTask
     conf = assigned.task
-    log.info('\t'.join((conf.owner.role,
-                        conf.jobName,
+    log.info('\t'.join((conf.job.role if conf.job else conf.owner.role,
+                        conf.job.name if conf.job else conf.jobName,
                         str(assigned.instanceId),
                         ScheduleStatus._VALUES_TO_NAMES[task.status],
                         assigned.taskId)))

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/client/commands/core.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/core.py b/src/main/python/apache/aurora/client/commands/core.py
index 58f419e..e634362 100644
--- a/src/main/python/apache/aurora/client/commands/core.py
+++ b/src/main/python/apache/aurora/client/commands/core.py
@@ -267,11 +267,12 @@ def diff(job_spec, config_file):
 
   pp = pprint.PrettyPrinter(indent=2)
   def pretty_print_task(task):
-    # The raw configuration is not interesting - we only care about what gets parsed.
+  # The raw configuration is not interesting - we only care about what gets parsed.
     task.configuration = None
     task.executorConfig = ExecutorConfig(
         name=AURORA_EXECUTOR_NAME,
         data=json.loads(task.executorConfig.data))
+
     return pp.pformat(vars(task))
 
   def pretty_print_tasks(tasks):
@@ -648,16 +649,17 @@ def status(args, options):
     return taskString
 
   def print_tasks(tasks):
-    for task in tasks:
-      taskString = print_task(task)
+    for scheduled in tasks:
+      taskString = print_task(scheduled)
 
+      assigned = scheduled.assignedTask
       log.info('role: %s, env: %s, name: %s, shard: %s, status: %s on %s\n%s' %
-             (task.assignedTask.task.owner.role,
-              task.assignedTask.task.environment,
-              task.assignedTask.task.jobName,
-              task.assignedTask.instanceId,
-              ScheduleStatus._VALUES_TO_NAMES[task.status],
-              task.assignedTask.slaveHost,
+             (assigned.task.job.role if assigned.task.job else assigned.task.owner.role,
+              assigned.task.job.environment if assigned.task.job else assigned.task.environment,
+              assigned.task.job.name if assigned.task.job else assigned.task.jobName,
+              assigned.instanceId,
+              ScheduleStatus._VALUES_TO_NAMES[scheduled.status],
+              assigned.slaveHost,
               taskString))
 
   api, job_key, _ = LiveJobDisambiguator.disambiguate_args_or_die(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/client/commands/ssh.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/ssh.py b/src/main/python/apache/aurora/client/commands/ssh.py
index d2b8bf6..b0916ed 100644
--- a/src/main/python/apache/aurora/client/commands/ssh.py
+++ b/src/main/python/apache/aurora/client/commands/ssh.py
@@ -82,8 +82,9 @@ def ssh(args, options):
 
   ssh_command = ['ssh', '-t']
 
-  role = first_task.assignedTask.task.owner.role
-  slave_host = first_task.assignedTask.slaveHost
+  assigned = first_task.assignedTask
+  role = assigned.task.job.role if assigned.task.job else assigned.task.owner.role
+  slave_host = assigned.slaveHost
 
   for tunnel in options.tunnels:
     try:
@@ -91,10 +92,10 @@ def ssh(args, options):
       port = int(port)
     except ValueError:
       die('Could not parse tunnel: %s.  Must be of form PORT:NAME' % tunnel)
-    if name not in first_task.assignedTask.assignedPorts:
-      die('Task %s has no port named %s' % (first_task.assignedTask.taskId, name))
+    if name not in assigned.assignedPorts:
+      die('Task %s has no port named %s' % (assigned.taskId, name))
     ssh_command += [
-        '-L', '%d:%s:%d' % (port, slave_host, first_task.assignedTask.assignedPorts[name])]
+        '-L', '%d:%s:%d' % (port, slave_host, assigned.assignedPorts[name])]
 
   ssh_command += ['%s@%s' % (options.ssh_user or role, slave_host), command]
   return subprocess.call(ssh_command)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/config/thrift.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/config/thrift.py b/src/main/python/apache/aurora/config/thrift.py
index 9ca806d..ba94ac3 100644
--- a/src/main/python/apache/aurora/config/thrift.py
+++ b/src/main/python/apache/aurora/config/thrift.py
@@ -196,6 +196,7 @@ def convert(job, metadata=frozenset(), ports=frozenset()):
     raise InvalidConfig('Task has invalid resources.  cpu/ramMb/diskMb must all be positive: '
         'cpu:%r ramMb:%r diskMb:%r' % (task.numCpus, task.ramMb, task.diskMb))
 
+  task.job = key
   task.owner = owner
   task.requestedPorts = ports
   task.taskLinks = not_empty_or(job.task_links(), {})

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/executor/aurora_executor.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/executor/aurora_executor.py b/src/main/python/apache/aurora/executor/aurora_executor.py
index 2c6423d..636b23d 100644
--- a/src/main/python/apache/aurora/executor/aurora_executor.py
+++ b/src/main/python/apache/aurora/executor/aurora_executor.py
@@ -38,7 +38,7 @@ class DefaultSandboxProvider(SandboxProvider):
   def from_assigned_task(self, assigned_task):
     return DirectorySandbox(
         os.path.realpath(self.SANDBOX_NAME),
-        assigned_task.task.owner.role)
+        assigned_task.task.job.role if assigned_task.task.job else assigned_task.task.owner.role)
 
 
 class AuroraExecutor(ExecutorBase, Observable):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/executor/common/announcer.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/executor/common/announcer.py b/src/main/python/apache/aurora/executor/common/announcer.py
index 74b2114..9e5bdc3 100644
--- a/src/main/python/apache/aurora/executor/common/announcer.py
+++ b/src/main/python/apache/aurora/executor/common/announcer.py
@@ -110,10 +110,11 @@ class DefaultAnnouncerCheckerProvider(AnnouncerCheckerProvider):
     return KazooClient(self.__ensemble, connection_retry=self.DEFAULT_RETRY_POLICY)
 
   def make_zk_path(self, assigned_task):
+    config = assigned_task.task
     role, environment, name = (
-        assigned_task.task.owner.role,
-        assigned_task.task.environment,
-        assigned_task.task.jobName)
+        config.job.role if config.job else config.owner.role,
+        config.job.environment if config.job else config.environment,
+        config.job.name if config.job else config.jobName)
     return posixpath.join(self.__root, role, environment, name)
 
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/main/python/apache/aurora/executor/thermos_task_runner.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/executor/thermos_task_runner.py b/src/main/python/apache/aurora/executor/thermos_task_runner.py
index bb99bd1..9a2faa0 100644
--- a/src/main/python/apache/aurora/executor/thermos_task_runner.py
+++ b/src/main/python/apache/aurora/executor/thermos_task_runner.py
@@ -367,7 +367,7 @@ class DefaultThermosTaskRunnerProvider(TaskRunnerProvider):
 
   def from_assigned_task(self, assigned_task, sandbox):
     task_id = assigned_task.taskId
-    role = assigned_task.task.owner.role
+    role = assigned_task.task.job.role if assigned_task.task.job else assigned_task.task.owner.role
     try:
       mesos_task = mesos_task_instance_from_assigned_task(assigned_task)
     except ValueError as e:

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/api/test_instance_watcher.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_instance_watcher.py b/src/test/python/apache/aurora/client/api/test_instance_watcher.py
index ae1b24b..abbbdbe 100644
--- a/src/test/python/apache/aurora/client/api/test_instance_watcher.py
+++ b/src/test/python/apache/aurora/client/api/test_instance_watcher.py
@@ -23,7 +23,6 @@ from apache.aurora.client.api.instance_watcher import InstanceWatcher
 from gen.apache.aurora.api.AuroraSchedulerManager import Client as scheduler_client
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
-    Identity,
     JobKey,
     Response,
     ResponseCode,
@@ -72,16 +71,13 @@ class InstanceWatcherTest(unittest.TestCase):
   EXPECTED_CYCLES = find_expected_cycles(WATCH_SECS, 3.0)
 
   def setUp(self):
-    self._role = 'mesos'
-    self._env = 'test'
-    self._name = 'jimbob'
     self._clock = FakeClock()
     self._event = FakeEvent(self._clock)
     self._scheduler = mox.MockObject(scheduler_client)
-    job_key = JobKey(name=self._name, environment=self._env, role=self._role)
+    self._job_key = JobKey(role='mesos', name='jimbob', environment='test')
     self._health_check = mox.MockObject(HealthCheck)
     self._watcher = InstanceWatcher(self._scheduler,
-                                 job_key,
+                                 self._job_key,
                                  self.RESTART_THRESHOLD,
                                  self.WATCH_SECS,
                                  health_check_interval_seconds=3,
@@ -90,9 +86,7 @@ class InstanceWatcherTest(unittest.TestCase):
 
   def get_tasks_status_query(self, instance_ids):
     query = TaskQuery()
-    query.owner = Identity(role=self._role)
-    query.environment = self._env
-    query.jobName = self._name
+    query.jobKeys = set([self._job_key])
     query.statuses = set([ScheduleStatus.RUNNING])
     query.instanceIds = set(instance_ids)
     return query
@@ -108,7 +102,7 @@ class InstanceWatcherTest(unittest.TestCase):
     response.result.scheduleStatusResult = ScheduleStatusResult(tasks=tasks)
 
     query = self.get_tasks_status_query(instance_ids)
-    for x in range(int(num_calls)):
+    for _ in range(int(num_calls)):
       self._scheduler.getTasksWithoutConfigs(query).AndReturn(response)
 
   def expect_io_error_in_get_statuses(self, instance_ids=WATCH_INSTANCES,
@@ -119,14 +113,14 @@ class InstanceWatcherTest(unittest.TestCase):
     response.result.scheduleStatusResult = ScheduleStatusResult(tasks=tasks)
 
     query = self.get_tasks_status_query(instance_ids)
-    for x in range(int(num_calls)):
+    for _ in range(int(num_calls)):
       self._scheduler.getTasksWithoutConfigs(query).AndRaise(IOError('oops'))
 
   def mock_health_check(self, task, status, retry):
     self._health_check.health(task).InAnyOrder().AndReturn((status, retry))
 
   def expect_health_check(self, instance, status, retry=True, num_calls=EXPECTED_CYCLES):
-    for x in range(int(num_calls)):
+    for _ in range(int(num_calls)):
       self.mock_health_check(self.create_task(instance), status, retry)
 
   def assert_watch_result(self, expected_failed_instances, instances_to_watch=WATCH_INSTANCES):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/api/test_sla.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_sla.py b/src/test/python/apache/aurora/client/api/test_sla.py
index 1117f24..50a6c47 100644
--- a/src/test/python/apache/aurora/client/api/test_sla.py
+++ b/src/test/python/apache/aurora/client/api/test_sla.py
@@ -27,6 +27,7 @@ from gen.apache.aurora.api.constants import LIVE_STATES
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
     Identity,
+    JobKey,
     Response,
     ResponseCode,
     Result,
@@ -67,6 +68,7 @@ class SlaTest(unittest.TestCase):
             slaveHost=host,
             task=TaskConfig(
                 production=prod if prod is not None else True,
+                job=JobKey(role=self._role, environment=self._env, name=name or self._name),
                 jobName=name or self._name,
                 owner=Identity(role=self._role),
                 environment=self._env)),
@@ -162,13 +164,9 @@ class SlaTest(unittest.TestCase):
     )
 
   def expect_task_status_call_job_scoped(self):
-    self._scheduler.getTasksWithoutConfigs.assert_called_once_with(
-        TaskQuery(
-            owner=Identity(role=self._role),
-            environment=self._env,
-            jobName=self._name,
-            statuses=LIVE_STATES)
-    )
+    self._scheduler.getTasksWithoutConfigs.assert_called_once_with(TaskQuery(
+        jobKeys=[self._job_key.to_thrift()],
+        statuses=LIVE_STATES))
 
   def expect_task_status_call_cluster_scoped(self):
     self._scheduler.getTasksWithoutConfigs.assert_called_with(TaskQuery(statuses=LIVE_STATES))

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/cli/test_api_from_cli.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_api_from_cli.py b/src/test/python/apache/aurora/client/cli/test_api_from_cli.py
index 95aa649..4de8bb4 100644
--- a/src/test/python/apache/aurora/client/cli/test_api_from_cli.py
+++ b/src/test/python/apache/aurora/client/cli/test_api_from_cli.py
@@ -23,19 +23,13 @@ from apache.aurora.client.cli.util import AuroraClientCommandTest
 
 from gen.apache.aurora.api import AuroraAdmin
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
     GetJobsResult,
-    Identity,
     JobConfiguration,
     JobKey,
     Response,
     ResponseCode,
     Result,
-    ScheduledTask,
-    ScheduleStatus,
     ScheduleStatusResult,
-    TaskConfig,
-    TaskEvent,
     TaskQuery
 )
 
@@ -46,35 +40,6 @@ class TestApiFromCLI(AuroraClientCommandTest):
   """
 
   @classmethod
-  def create_mock_scheduled_tasks(cls):
-    tasks = []
-    for name in ['foo', 'bar', 'baz']:
-      task = Mock(spec=ScheduledTask)
-      task.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      task.failure_count = 0
-      task.assignedTask = Mock(spec=AssignedTask)
-      task.assignedTask.slaveHost = 'slavehost'
-      task.assignedTask.task = Mock(spec=TaskConfig)
-      task.assignedTask.task.maxTaskFailures = 1
-      task.assignedTask.task.metadata = []
-      task.assignedTask.task.owner = Identity(role='bozo')
-      task.assignedTask.task.environment = 'test'
-      task.assignedTask.task.jobName = 'woops'
-      task.assignedTask.task.numCpus = 2
-      task.assignedTask.task.ramMb = 2
-      task.assignedTask.task.diskMb = 2
-      task.assignedTask.instanceId = 4237894
-      task.assignedTask.assignedPorts = None
-      task.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
-      mockEvent.timestamp = 28234726395
-      mockEvent.status = ScheduleStatus.RUNNING
-      mockEvent.message = "Hi there"
-      task.taskEvents = [mockEvent]
-      tasks.append(task)
-    return tasks
-
-  @classmethod
   def create_mock_scheduled_task_no_metadata(cls):
     result = cls.create_mock_scheduled_tasks()
     for task in result:
@@ -104,7 +69,7 @@ class TestApiFromCLI(AuroraClientCommandTest):
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
     resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
-    resp.result.scheduleStatusResult.tasks = set(cls.create_mock_scheduled_tasks())
+    resp.result.scheduleStatusResult.tasks = set(cls.create_scheduled_tasks())
     return resp
 
   @classmethod

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/cli/test_diff.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_diff.py b/src/test/python/apache/aurora/client/cli/test_diff.py
index 1081769..78694d7 100644
--- a/src/test/python/apache/aurora/client/cli/test_diff.py
+++ b/src/test/python/apache/aurora/client/cli/test_diff.py
@@ -24,17 +24,11 @@ from apache.aurora.client.cli.util import AuroraClientCommandTest
 
 from gen.apache.aurora.api.constants import ACTIVE_STATES
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
-    ExecutorConfig,
-    Identity,
     JobConfiguration,
     JobKey,
     PopulateJobResult,
     ResponseCode,
-    ScheduleStatus,
     ScheduleStatusResult,
-    TaskConfig,
-    TaskEvent,
     TaskQuery
 )
 
@@ -44,7 +38,6 @@ class TestDiffCommand(AuroraClientCommandTest):
   def setup_mock_options(cls):
     """set up to get a mock options object."""
     mock_options = Mock()
-    mock_options = Mock()
     mock_options.env = None
     mock_options.json = False
     mock_options.bindings = {}
@@ -54,41 +47,10 @@ class TestDiffCommand(AuroraClientCommandTest):
     return mock_options
 
   @classmethod
-  def create_mock_scheduled_tasks(cls):
-    jobs = []
-    for name in ['foo', 'bar', 'baz']:
-      job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      job.failure_count = 0
-      job.assignedTask = Mock(spec=AssignedTask)
-      job.assignedTask.slaveHost = 'slavehost'
-      job.assignedTask.task = Mock(spec=TaskConfig)
-      job.assignedTask.task.maxTaskFailures = 1
-      job.assignedTask.task.executorConfig = Mock(spec=ExecutorConfig)
-      job.assignedTask.task.executorConfig.data = Mock()
-      job.assignedTask.task.metadata = []
-      job.assignedTask.task.owner = Identity(role='bozo')
-      job.assignedTask.task.environment = 'test'
-      job.assignedTask.task.jobName = 'woops'
-      job.assignedTask.task.numCpus = 2
-      job.assignedTask.task.ramMb = 2
-      job.assignedTask.task.diskMb = 2
-      job.assignedTask.instanceId = 4237894
-      job.assignedTask.assignedPorts = None
-      job.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
-      mockEvent.timestamp = 28234726395
-      mockEvent.status = ScheduleStatus.RUNNING
-      mockEvent.message = "Hi there"
-      job.taskEvents = [mockEvent]
-      jobs.append(job)
-    return jobs
-
-  @classmethod
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
     resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
-    resp.result.scheduleStatusResult.tasks = set(cls.create_mock_scheduled_tasks())
+    resp.result.scheduleStatusResult.tasks = set(cls.create_scheduled_tasks())
     return resp
 
   @classmethod
@@ -100,7 +62,8 @@ class TestDiffCommand(AuroraClientCommandTest):
     populate = cls.create_simple_success_response()
     populate.result.populateJobResult = Mock(spec=PopulateJobResult)
     api.populateJobConfig.return_value = populate
-    populate.result.populateJobResult.populatedDEPRECATED = cls.create_mock_scheduled_tasks()
+    tasks = set(task.assignedTask.task for task in cls.create_scheduled_tasks())
+    populate.result.populateJobResult.populatedDEPRECATED = tasks
     return populate
 
   def test_successful_diff(self):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/cli/test_status.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_status.py b/src/test/python/apache/aurora/client/cli/test_status.py
index 49ac2a4..8894a1e 100644
--- a/src/test/python/apache/aurora/client/cli/test_status.py
+++ b/src/test/python/apache/aurora/client/cli/test_status.py
@@ -43,45 +43,22 @@ from gen.apache.aurora.api.ttypes import (
 
 class TestJobStatus(AuroraClientCommandTest):
   @classmethod
-  def create_mock_scheduled_tasks(cls):
-    tasks = []
+  def create_scheduled_tasks(cls):
+    tasks = AuroraClientCommandTest.create_scheduled_tasks()
     instance = 0
-    for name in ['foo', 'bar', 'baz']:
+    for task in tasks:
       instance += 1
-      event = TaskEvent(
-        timestamp=28234726395,
-        status=ScheduleStatus.RUNNING,
-        message="Hi there"
-      )
-      task = ScheduledTask(
-        failureCount=0,
-        assignedTask=AssignedTask(
-          slaveHost='slavehost',
-          task=TaskConfig(
-            maxTaskFailures=1,
-            metadata=[],
-            owner=Identity(role='bozo'),
-            environment='test',
-            jobName='woops',
-            numCpus=2,
-            ramMb=2,
-            diskMb=2
-          ),
-          instanceId=instance,
-          assignedPorts=None,
-        ),
-        status=ScheduleStatus.RUNNING,
-        taskEvents=[event]
-      )
-      tasks.append(task)
+      task.assignedTask.instanceId = instance
+      task.assignedTask.task.job = JobKey(role='bozo', environment='test', name='woops')
+      task.assignedTask.task.jobName = 'woops'
     return tasks
 
   @classmethod
-  def create_mock_inactive_tasks(cls):
-    jobs = []
+  def create_inactive_tasks(cls):
     instance = 0
     INACTIVE_STATUSES = [ScheduleStatus.KILLED, ScheduleStatus.FINISHED, ScheduleStatus.FAILED]
-    for instance in range(3):
+    tasks = cls.create_scheduled_tasks()
+    for task in tasks:
       events = []
       for i in range(3):
         event = TaskEvent(
@@ -89,36 +66,22 @@ class TestJobStatus(AuroraClientCommandTest):
           status=INACTIVE_STATUSES[i],
           message="Hi there")
         events.append(event)
-      job = ScheduledTask(
-        failureCount=3,
-        assignedTask=AssignedTask(
-          slaveHost='slavehost',
-          task=TaskConfig(
-            maxTaskFailures=1,
-            metadata=[],
-            owner=Identity(role='bozo'),
-            environment='test',
-            jobName='woops',
-            numCpus=2,
-            ramMb=2,
-            diskMb=2),
-          instanceId=instance,
-          assignedPorts=None),
-        status=INACTIVE_STATUSES[instance],
-        taskEvents=events)
-      jobs.append(job)
-    return set(jobs)
+      task.taskEvents = events
+      task.status = INACTIVE_STATUSES[instance]
+      task.assignedTask.instanceId = instance
+      instance += 1
+    return set(tasks)
 
   @classmethod
   def create_mock_scheduled_task_no_metadata(cls):
-    result = cls.create_mock_scheduled_tasks()
+    result = cls.create_scheduled_tasks()
     for job in result:
       job.assignedTask.task.metadata = None
     return result
 
   @classmethod
   def create_mock_scheduled_task_with_metadata(cls):
-    result = cls.create_mock_scheduled_tasks()
+    result = cls.create_scheduled_tasks()
     for job in result:
       job.assignedTask.task.metadata = [Metadata("meta", "data"), Metadata("data", "meta")]
     return result
@@ -144,7 +107,7 @@ class TestJobStatus(AuroraClientCommandTest):
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
     resp.result.scheduleStatusResult = ScheduleStatusResult(
-      tasks=set(cls.create_mock_scheduled_tasks()))
+      tasks=set(cls.create_scheduled_tasks()))
     return resp
 
   @classmethod
@@ -157,7 +120,7 @@ class TestJobStatus(AuroraClientCommandTest):
   @classmethod
   def create_status_with_inactives(cls):
     resp = cls.create_status_null_metadata()
-    resp.result.scheduleStatusResult.tasks |= cls.create_mock_inactive_tasks()
+    resp.result.scheduleStatusResult.tasks |= cls.create_inactive_tasks()
     return resp
 
   @classmethod
@@ -182,6 +145,7 @@ class TestJobStatus(AuroraClientCommandTest):
           slaveId="random_machine_id",
           slaveHost="junk.nothing",
           task=TaskConfig(
+            job=JobKey(role="nobody", environment="prod", name='flibber'),
             owner=Identity(role="nobody"),
             environment="prod",
             jobName="flibber",
@@ -496,6 +460,11 @@ class TestJobStatus(AuroraClientCommandTest):
                     "owner": {
                       "role": "nobody"
                     },
+                    "job": {
+                      "environment": "prod",
+                      "role": "nobody",
+                      "name": "flibber"
+                    },
                     "production": false,
                     "diskMb": 4096,
                     "ramMb": 2048,
@@ -544,6 +513,11 @@ class TestJobStatus(AuroraClientCommandTest):
                     "owner": {
                       "role": "nobody"
                     },
+                    "job": {
+                      "environment": "prod",
+                      "role": "nobody",
+                      "name": "flibber"
+                    },
                     "production": false,
                     "diskMb": 4096,
                     "ramMb": 2048,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/cli/test_task_run.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_task_run.py b/src/test/python/apache/aurora/client/cli/test_task_run.py
index 16fde14..12163df 100644
--- a/src/test/python/apache/aurora/client/cli/test_task_run.py
+++ b/src/test/python/apache/aurora/client/cli/test_task_run.py
@@ -21,14 +21,10 @@ from apache.aurora.client.cli.client import AuroraCommandLine
 from apache.aurora.client.cli.util import AuroraClientCommandTest
 
 from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
-    Identity,
     JobKey,
     ResponseCode,
     ScheduleStatus,
     ScheduleStatusResult,
-    TaskConfig,
-    TaskEvent,
     TaskQuery
 )
 
@@ -42,41 +38,10 @@ def mock_log(level, msg):
 class TestRunCommand(AuroraClientCommandTest):
 
   @classmethod
-  def create_mock_scheduled_tasks(cls):
-    jobs = []
-    for name in ['foo', 'bar', 'baz']:
-      job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      job.failure_count = 0
-      job.assignedTask = Mock(spec=AssignedTask)
-      job.assignedTask.taskId = 1287391823
-      job.assignedTask.slaveHost = 'slavehost'
-      job.assignedTask.task = Mock(spec=TaskConfig)
-      job.assignedTask.task.executorConfig = Mock()
-      job.assignedTask.task.maxTaskFailures = 1
-      job.assignedTask.task.metadata = []
-      job.assignedTask.task.owner = Identity(role='bozo')
-      job.assignedTask.task.environment = 'test'
-      job.assignedTask.task.jobName = 'woops'
-      job.assignedTask.task.numCpus = 2
-      job.assignedTask.task.ramMb = 2
-      job.assignedTask.task.diskMb = 2
-      job.assignedTask.instanceId = 4237894
-      job.assignedTask.assignedPorts = {}
-      job.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
-      mockEvent.timestamp = 28234726395
-      mockEvent.status = ScheduleStatus.RUNNING
-      mockEvent.message = "Hi there"
-      job.taskEvents = [mockEvent]
-      jobs.append(job)
-    return jobs
-
-  @classmethod
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
     resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
-    resp.result.scheduleStatusResult.tasks = cls.create_mock_scheduled_tasks()
+    resp.result.scheduleStatusResult.tasks = cls.create_scheduled_tasks()
     return resp
 
   @classmethod
@@ -146,42 +111,12 @@ class TestRunCommand(AuroraClientCommandTest):
 
 
 class TestSshCommand(AuroraClientCommandTest):
-  @classmethod
-  def create_mock_scheduled_tasks(cls):
-    jobs = []
-    for name in ['foo', 'bar', 'baz']:
-      job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      job.failure_count = 0
-      job.assignedTask = Mock(spec=AssignedTask)
-      job.assignedTask.taskId = 1287391823
-      job.assignedTask.slaveHost = 'slavehost'
-      job.assignedTask.task = Mock(spec=TaskConfig)
-      job.assignedTask.task.executorConfig = Mock()
-      job.assignedTask.task.maxTaskFailures = 1
-      job.assignedTask.task.metadata = []
-      job.assignedTask.task.owner = Identity(role='bozo')
-      job.assignedTask.task.environment = 'test'
-      job.assignedTask.task.jobName = 'woops'
-      job.assignedTask.task.numCpus = 2
-      job.assignedTask.task.ramMb = 2
-      job.assignedTask.task.diskMb = 2
-      job.assignedTask.instanceId = 4237894
-      job.assignedTask.assignedPorts = {}
-      job.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
-      mockEvent.timestamp = 28234726395
-      mockEvent.status = ScheduleStatus.RUNNING
-      mockEvent.message = "Hi there"
-      job.taskEvents = [mockEvent]
-      jobs.append(job)
-    return jobs
 
   @classmethod
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
     resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
-    resp.result.scheduleStatusResult.tasks = cls.create_mock_scheduled_tasks()
+    resp.result.scheduleStatusResult.tasks = cls.create_scheduled_tasks()
     return resp
 
   @classmethod

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/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 1ec5483..08f277e 100644
--- a/src/test/python/apache/aurora/client/cli/test_update.py
+++ b/src/test/python/apache/aurora/client/cli/test_update.py
@@ -120,7 +120,11 @@ class TestUpdateCommand(AuroraClientCommandTest):
     populate = cls.create_simple_success_response()
     populate.result.populateJobResult = Mock(spec=PopulateJobResult)
     api.populateJobConfig.return_value = populate
-    configs = [TaskConfig(numCpus=1.0, ramMb=1, diskMb=1) for i in range(count)]
+    configs = [TaskConfig(
+        numCpus=1.0,
+        ramMb=1,
+        diskMb=1,
+        job=JobKey(role='bozo', environment='test', name='hello')) for i in range(count)]
     populate.result.populateJobResult.populatedDEPRECATED = set(configs)
     return populate
 
@@ -147,7 +151,12 @@ class TestUpdateCommand(AuroraClientCommandTest):
     scheduler.getTasksWithoutConfigs.return_value = status_response
     schedule_status = Mock(spec=ScheduleStatusResult)
     status_response.result.scheduleStatusResult = schedule_status
-    task_config = TaskConfig(numCpus=1.0, ramMb=10, diskMb=1)
+    task_config = TaskConfig(
+        numCpus=1.0,
+        ramMb=10,
+        diskMb=1,
+        job=JobKey(role='bozo', environment='test', name='hello'))
+
     # This should be a list of ScheduledTask's.
     schedule_status.tasks = []
     for i in range(20):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/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 3fa609a..796c4f9 100644
--- a/src/test/python/apache/aurora/client/cli/util.py
+++ b/src/test/python/apache/aurora/client/cli/util.py
@@ -25,6 +25,9 @@ from apache.aurora.common.clusters import Clusters
 
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
+    ExecutorConfig,
+    Identity,
+    JobKey,
     Response,
     ResponseCode,
     Result,
@@ -126,7 +129,6 @@ class AuroraClientCommandTest(unittest.TestCase):
   @classmethod
   def create_mock_api(cls):
     """Builds up a mock API object, with a mock SchedulerProxy"""
-    mock_api = Mock(spec=HookedAuroraClientAPI)
     mock_scheduler = Mock()
     mock_scheduler.url = "http://something_or_other"
     mock_scheduler_client = Mock()
@@ -175,6 +177,38 @@ class AuroraClientCommandTest(unittest.TestCase):
     return mock_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 = Mock()
+      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]
+      tasks.append(task)
+    return tasks
+
+  @classmethod
   def setup_get_tasks_status_calls(cls, scheduler):
     status_response = cls.create_status_call_result()
     scheduler.getTasksWithoutConfigs.return_value = status_response

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/commands/test_diff.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_diff.py b/src/test/python/apache/aurora/client/commands/test_diff.py
index c8d0145..9f1d459 100644
--- a/src/test/python/apache/aurora/client/commands/test_diff.py
+++ b/src/test/python/apache/aurora/client/commands/test_diff.py
@@ -31,6 +31,7 @@ from gen.apache.aurora.api.ttypes import (
     JobKey,
     PopulateJobResult,
     ResponseCode,
+    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
     TaskConfig,
@@ -44,7 +45,6 @@ class TestDiffCommand(AuroraClientCommandTest):
   def setup_mock_options(cls):
     """set up to get a mock options object."""
     mock_options = Mock()
-    mock_options = Mock()
     mock_options.env = None
     mock_options.json = False
     mock_options.bindings = {}
@@ -58,8 +58,7 @@ class TestDiffCommand(AuroraClientCommandTest):
   def create_mock_scheduled_tasks(cls):
     jobs = []
     for name in ['foo', 'bar', 'baz']:
-      job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
+      job = Mock(spec=ScheduledTask)
       job.failure_count = 0
       job.assignedTask = Mock(spec=AssignedTask)
       job.assignedTask.slaveHost = 'slavehost'
@@ -68,6 +67,7 @@ class TestDiffCommand(AuroraClientCommandTest):
       job.assignedTask.task.executorConfig = Mock(spec=ExecutorConfig)
       job.assignedTask.task.executorConfig.data = Mock()
       job.assignedTask.task.metadata = []
+      job.assignedTask.task.job = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
       job.assignedTask.task.owner = Identity(role='mchucarroll')
       job.assignedTask.task.environment = 'test'
       job.assignedTask.task.jobName = 'woops'
@@ -101,7 +101,8 @@ class TestDiffCommand(AuroraClientCommandTest):
     populate = cls.create_simple_success_response()
     populate.result.populateJobResult = Mock(spec=PopulateJobResult)
     api.populateJobConfig.return_value = populate
-    populate.result.populateJobResult.populatedDEPRECATED = cls.create_mock_scheduled_tasks()
+    tasks = set(task.assignedTask.task for task in cls.create_mock_scheduled_tasks())
+    populate.result.populateJobResult.populatedDEPRECATED = tasks
     return populate
 
   def test_successful_diff(self):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/commands/test_ssh.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_ssh.py b/src/test/python/apache/aurora/client/commands/test_ssh.py
index abb657b..cf9f425 100644
--- a/src/test/python/apache/aurora/client/commands/test_ssh.py
+++ b/src/test/python/apache/aurora/client/commands/test_ssh.py
@@ -49,8 +49,9 @@ class TestSshCommand(AuroraClientCommandTest):
   def create_mock_scheduled_tasks(cls):
     jobs = []
     for name in ['foo', 'bar', 'baz']:
+      job_key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
       job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
+      job.key = job_key
       job.failure_count = 0
       job.assignedTask = Mock(spec=AssignedTask)
       job.assignedTask.taskId = 1287391823
@@ -59,9 +60,10 @@ class TestSshCommand(AuroraClientCommandTest):
       job.assignedTask.task.executorConfig = Mock()
       job.assignedTask.task.maxTaskFailures = 1
       job.assignedTask.task.metadata = []
-      job.assignedTask.task.owner = Identity(role='mchucarroll')
-      job.assignedTask.task.environment = 'test'
-      job.assignedTask.task.jobName = 'woops'
+      job.assignedTask.task.job = job_key
+      job.assignedTask.task.owner = Identity(role=cls.TEST_ROLE)
+      job.assignedTask.task.environment = cls.TEST_ENV
+      job.assignedTask.task.jobName = name
       job.assignedTask.task.numCpus = 2
       job.assignedTask.task.ramMb = 2
       job.assignedTask.task.diskMb = 2

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/commands/test_status.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_status.py b/src/test/python/apache/aurora/client/commands/test_status.py
index 6397635..9eb8def 100644
--- a/src/test/python/apache/aurora/client/commands/test_status.py
+++ b/src/test/python/apache/aurora/client/commands/test_status.py
@@ -46,17 +46,19 @@ class TestListJobs(AuroraClientCommandTest):
   def create_mock_scheduled_tasks(cls):
     jobs = []
     for name in ['foo', 'bar', 'baz']:
+      job_key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
       job = Mock()
-      job.key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
+      job.key = job_key
       job.failure_count = 0
       job.assignedTask = Mock(spec=AssignedTask)
       job.assignedTask.slaveHost = 'slavehost'
       job.assignedTask.task = Mock(spec=TaskConfig)
       job.assignedTask.task.maxTaskFailures = 1
       job.assignedTask.task.metadata = []
-      job.assignedTask.task.owner = Identity(role='mchucarroll')
-      job.assignedTask.task.environment = 'test'
-      job.assignedTask.task.jobName = 'woops'
+      job.assignedTask.task.job = job_key
+      job.assignedTask.task.owner = Identity(role=cls.TEST_ROLE)
+      job.assignedTask.task.environment = cls.TEST_ENV
+      job.assignedTask.task.jobName = name
       job.assignedTask.task.numCpus = 2
       job.assignedTask.task.ramMb = 2
       job.assignedTask.task.diskMb = 2

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/client/commands/test_update.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_update.py b/src/test/python/apache/aurora/client/commands/test_update.py
index 07cbe53..555ea0d 100644
--- a/src/test/python/apache/aurora/client/commands/test_update.py
+++ b/src/test/python/apache/aurora/client/commands/test_update.py
@@ -148,7 +148,12 @@ class TestUpdateCommand(AuroraClientCommandTest):
     api.populateJobConfig.return_value = populate
     configs = []
     for _ in range(20):
-      task_config = TaskConfig(numCpus=1.0, ramMb=1, diskMb=1)
+      task_config = TaskConfig(
+          numCpus=1.0,
+          ramMb=1,
+          diskMb=1,
+          job=JobKey(role='mchucarroll', environment='test', name='hello'))
+
       configs.append(task_config)
     populate.result.populateJobResult.populatedDEPRECATED = set(configs)
     return populate
@@ -176,7 +181,12 @@ class TestUpdateCommand(AuroraClientCommandTest):
     scheduler_proxy.getTasksWithoutConfigs.return_value = status_response
     schedule_status = Mock(spec=ScheduleStatusResult)
     status_response.result.scheduleStatusResult = schedule_status
-    task_config = TaskConfig(numCpus=1.0, ramMb=10, diskMb=1)
+    task_config = TaskConfig(
+        numCpus=1.0,
+        ramMb=10,
+        diskMb=1,
+        job=JobKey(role='mchucarroll', environment='test', name='hello'))
+
     # This should be a list of ScheduledTask's.
     schedule_status.tasks = []
     for i in range(20):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/config/test_thrift.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/config/test_thrift.py b/src/test/python/apache/aurora/config/test_thrift.py
index 1dd9e79..654c0b5 100644
--- a/src/test/python/apache/aurora/config/test_thrift.py
+++ b/src/test/python/apache/aurora/config/test_thrift.py
@@ -44,14 +44,16 @@ HELLO_WORLD = Job(
 
 def test_simple_config():
   job = convert_pystachio_to_thrift(HELLO_WORLD)
+  expected_key = JobKey(
+      role=HELLO_WORLD.role().get(),
+      environment=HELLO_WORLD.environment().get(),
+      name=HELLO_WORLD.name().get())
   assert job.instanceCount == 1
   tti = job.taskConfig
-  assert job.key == JobKey(
-    role=HELLO_WORLD.role().get(),
-    environment=HELLO_WORLD.environment().get(),
-    name=HELLO_WORLD.name().get())
+  assert job.key == expected_key
   assert job.owner == Identity(role=HELLO_WORLD.role().get(), user=getpass.getuser())
   assert job.cronSchedule is None
+  assert tti.job == expected_key
   assert tti.jobName == 'hello_world'
   assert tti.isService is False
   assert tti.numCpus == 0.1

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/executor/common/test_announcer.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/executor/common/test_announcer.py b/src/test/python/apache/aurora/executor/common/test_announcer.py
index 5694335..a4ab532 100644
--- a/src/test/python/apache/aurora/executor/common/test_announcer.py
+++ b/src/test/python/apache/aurora/executor/common/test_announcer.py
@@ -188,11 +188,21 @@ def test_announcer_under_abnormal_circumstances():
 
 def make_assigned_task(thermos_config, assigned_ports=None):
   from gen.apache.aurora.api.constants import AURORA_EXECUTOR_NAME
-  from gen.apache.aurora.api.ttypes import AssignedTask, ExecutorConfig, Identity, TaskConfig
+  from gen.apache.aurora.api.ttypes import (
+      AssignedTask,
+      ExecutorConfig,
+      Identity,
+      JobKey,
+      TaskConfig
+  )
 
   assigned_ports = assigned_ports or {}
   executor_config = ExecutorConfig(name=AURORA_EXECUTOR_NAME, data=thermos_config.json_dumps())
   task_config = TaskConfig(
+      job=JobKey(
+          role=thermos_config.role().get(),
+          environment="prod",
+          name=thermos_config.name().get()),
       owner=Identity(role=thermos_config.role().get(), user=thermos_config.role().get()),
       environment=thermos_config.environment().get(),
       jobName=thermos_config.name().get(),

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/06935c04/src/test/python/apache/aurora/executor/test_thermos_executor.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/executor/test_thermos_executor.py b/src/test/python/apache/aurora/executor/test_thermos_executor.py
index 65e8cce..16a4011 100644
--- a/src/test/python/apache/aurora/executor/test_thermos_executor.py
+++ b/src/test/python/apache/aurora/executor/test_thermos_executor.py
@@ -55,7 +55,7 @@ from apache.thermos.core.runner import TaskRunner
 from apache.thermos.monitoring.monitor import TaskMonitor
 
 from gen.apache.aurora.api.constants import AURORA_EXECUTOR_NAME
-from gen.apache.aurora.api.ttypes import AssignedTask, ExecutorConfig, Identity, TaskConfig
+from gen.apache.aurora.api.ttypes import AssignedTask, ExecutorConfig, Identity, JobKey, TaskConfig
 
 if 'THERMOS_DEBUG' in os.environ:
   LogOptions.set_stderr_log_level('google:DEBUG')
@@ -137,6 +137,7 @@ def make_task(thermos_config, assigned_ports={}, **kw):
           executorConfig=ExecutorConfig(
               name=AURORA_EXECUTOR_NAME,
               data=thermos_config.json_dumps()),
+          job=JobKey(role=role, environment='env', name='name'),
           owner=Identity(role=role, user=role)),
       assignedPorts=assigned_ports,
       **kw)