aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From wfar...@apache.org
Subject [1/2] Convert most uses of Mock to create_autospec, remove some uses of mocking altogether.
Date Thu, 06 Nov 2014 06:00:10 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master 0ad2fcabb -> d01f0c382


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_create.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_create.py b/src/test/python/apache/aurora/client/commands/test_create.py
index be69ed5..d7c32b3 100644
--- a/src/test/python/apache/aurora/client/commands/test_create.py
+++ b/src/test/python/apache/aurora/client/commands/test_create.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import create_autospec, patch
 from pystachio.config import Config
 from twitter.common.contextutil import temporary_file
 
@@ -38,23 +38,24 @@ class TestClientCreateCommand(AuroraClientCommandTest):
   @classmethod
   def setup_mock_options(cls):
     """set up to get a mock options object."""
-    mock_options = Mock()
+    mock_options = create_autospec(spec={}, instance=True)
     mock_options.json = False
     mock_options.bindings = {}
     mock_options.open_browser = False
     mock_options.cluster = None
     mock_options.wait_until = 'RUNNING'  # or 'FINISHED' for other tests
     mock_options.disable_all_hooks = False
+    mock_options.disable_all_hooks_reason = None
     return mock_options
 
   @classmethod
   def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    mock_task = Mock(spec=ScheduledTask)
-    mock_task.assignedTask = Mock(spec=AssignedTask)
+    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 = Mock(spec=TaskEvent)
+    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
     mock_task_event.timestamp = initial_time
     mock_task.taskEvents = [mock_task_event]
     return mock_task
@@ -62,7 +63,10 @@ class TestClientCreateCommand(AuroraClientCommandTest):
   @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
     mock_query_result = cls.create_simple_success_response()
-    mock_query_result.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
+    mock_query_result.result.scheduleStatusResult = create_autospec(
+        spec=ScheduleStatusResult,
+        spec_set=False,
+        instance=True)
     mock_task_one = cls.create_mock_task('hello', 0, 1000, scheduleStatus)
     mock_task_two = cls.create_mock_task('hello', 1, 1004, scheduleStatus)
     mock_query_result.result.scheduleStatusResult.tasks = [mock_task_one, mock_task_two]

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/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 9a486c3..b48b53d 100644
--- a/src/test/python/apache/aurora/client/commands/test_diff.py
+++ b/src/test/python/apache/aurora/client/commands/test_diff.py
@@ -15,7 +15,7 @@
 import contextlib
 import os
 
-from mock import Mock, patch
+from mock import create_autospec, Mock, patch
 from pystachio.config import Config
 from twitter.common.contextutil import temporary_file
 
@@ -32,6 +32,7 @@ from gen.apache.aurora.api.ttypes import (
     JobKey,
     PopulateJobResult,
     ResponseCode,
+    Result,
     ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
@@ -59,14 +60,13 @@ class TestDiffCommand(AuroraClientCommandTest):
   def create_mock_scheduled_tasks(cls):
     jobs = []
     for name in ['foo', 'bar', 'baz']:
-      job = Mock(spec=ScheduledTask)
+      job = create_autospec(spec=ScheduledTask, instance=True)
       job.failure_count = 0
-      job.assignedTask = Mock(spec=AssignedTask)
+      job.assignedTask = create_autospec(spec=AssignedTask, instance=True)
       job.assignedTask.slaveHost = 'slavehost'
-      job.assignedTask.task = Mock(spec=TaskConfig)
+      job.assignedTask.task = create_autospec(spec=TaskConfig, instance=True)
       job.assignedTask.task.maxTaskFailures = 1
-      job.assignedTask.task.executorConfig = Mock(spec=ExecutorConfig)
-      job.assignedTask.task.executorConfig.data = Mock()
+      job.assignedTask.task.executorConfig = ExecutorConfig(name='name', data='fake data')
       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')
@@ -78,7 +78,7 @@ class TestDiffCommand(AuroraClientCommandTest):
       job.assignedTask.instanceId = 4237894
       job.assignedTask.assignedPorts = None
       job.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
+      mockEvent = create_autospec(spec=TaskEvent, instance=True)
       mockEvent.timestamp = 28234726395
       mockEvent.status = ScheduleStatus.RUNNING
       mockEvent.message = "Hi there"
@@ -89,8 +89,8 @@ class TestDiffCommand(AuroraClientCommandTest):
   @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 = Result(
+        scheduleStatusResult=ScheduleStatusResult(tasks=set(cls.create_mock_scheduled_tasks())))
     return resp
 
   @classmethod
@@ -100,10 +100,12 @@ class TestDiffCommand(AuroraClientCommandTest):
   @classmethod
   def setup_populate_job_config(cls, api):
     populate = cls.create_simple_success_response()
-    populate.result.populateJobResult = Mock(spec=PopulateJobResult)
+
     api.populateJobConfig.return_value = populate
     tasks = set(task.assignedTask.task for task in cls.create_mock_scheduled_tasks())
-    populate.result.populateJobResult.populatedDEPRECATED = tasks
+    populate.result = Result(populateJobResult=PopulateJobResult(
+      populatedDEPRECATED=tasks
+    ))
     return populate
 
   def test_successful_diff(self):

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_hooks.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_hooks.py b/src/test/python/apache/aurora/client/commands/test_hooks.py
index 638fba6..de1781e 100644
--- a/src/test/python/apache/aurora/client/commands/test_hooks.py
+++ b/src/test/python/apache/aurora/client/commands/test_hooks.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import create_autospec, patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.commands.core import create
@@ -25,6 +25,7 @@ from .util import AuroraClientCommandTest
 
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
+    Result,
     ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
@@ -50,7 +51,7 @@ class TestClientCreateCommand(AuroraClientCommandTest):
   @classmethod
   def setup_mock_options(cls):
     """set up to get a mock options object."""
-    mock_options = Mock()
+    mock_options = create_autospec(spec={}, instance=True)
     mock_options.json = False
     mock_options.bindings = {}
     mock_options.open_browser = False
@@ -61,12 +62,12 @@ class TestClientCreateCommand(AuroraClientCommandTest):
 
   @classmethod
   def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    mock_task = Mock(spec=ScheduledTask)
-    mock_task.assignedTask = Mock(spec=AssignedTask)
+    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 = Mock(spec=TaskEvent)
+    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
     mock_task_event.timestamp = initial_time
     mock_task.taskEvents = [mock_task_event]
     return mock_task
@@ -74,14 +75,15 @@ class TestClientCreateCommand(AuroraClientCommandTest):
   @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
     mock_query_result = cls.create_simple_success_response()
-    mock_query_result.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
+
     if scheduleStatus == ScheduleStatus.INIT:
       # status query result for before job is launched.
-      mock_query_result.result.scheduleStatusResult.tasks = []
+      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_query_result.result.scheduleStatusResult.tasks = [mock_task_one, mock_task_two]
+      tasks = [mock_task_one, mock_task_two]
+    mock_query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
     return mock_query_result
 
   @classmethod

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_kill.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_kill.py b/src/test/python/apache/aurora/client/commands/test_kill.py
index 8194c96..4ac742f 100644
--- a/src/test/python/apache/aurora/client/commands/test_kill.py
+++ b/src/test/python/apache/aurora/client/commands/test_kill.py
@@ -13,8 +13,9 @@
 #
 
 import contextlib
+import threading
 
-from mock import Mock, patch
+from mock import create_autospec, Mock, patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.commands.core import CoreCommandHook, kill, killall
@@ -31,6 +32,7 @@ from gen.apache.aurora.api.ttypes import (
     ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
+    TaskConfig,
     TaskEvent,
     TaskQuery
 )
@@ -41,7 +43,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
   @classmethod
   def setup_mock_options(cls):
     """set up to get a mock options object."""
-    mock_options = Mock()
+    mock_options = create_autospec(spec=[], instance=True)
     mock_options.open_browser = False
     mock_options.shards = None
     mock_options.cluster = None
@@ -49,12 +51,14 @@ class TestClientKillCommand(AuroraClientCommandTest):
     mock_options.batch_size = None
     mock_options.max_total_failures = 1
     mock_options.disable_all_hooks = False
+    mock_options.disable_all_hooks_reason = None
+    mock_options.max_failures_option = None
     return mock_options
 
   @classmethod
   def setup_mock_api_factory(cls):
     mock_api_factory, mock_api = cls.create_mock_api_factory()
-    mock_api_factory.return_value.kill_job.return_value = cls.get_kill_job_response()
+    mock_api_factory('x').return_value.kill_job.return_value = cls.get_kill_job_response()
     return mock_api_factory
 
   @classmethod
@@ -71,12 +75,12 @@ class TestClientKillCommand(AuroraClientCommandTest):
 
   @classmethod
   def create_mock_task(cls, task_id, instance_id, initial_time, status):
-    mock_task = Mock(spec=ScheduledTask)
-    mock_task.assignedTask = Mock(spec=AssignedTask)
+    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 = Mock(spec=TaskEvent)
+    mock_task_event = create_autospec(spec=TaskEvent, instance=True)
     mock_task_event.timestamp = initial_time
     mock_task.taskEvents = [mock_task_event]
     return mock_task
@@ -84,7 +88,10 @@ class TestClientKillCommand(AuroraClientCommandTest):
   @classmethod
   def create_mock_status_query_result(cls, scheduleStatus):
     mock_query_result = cls.create_simple_success_response()
-    mock_query_result.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
+    mock_query_result.result.scheduleStatusResult = create_autospec(
+        spec=ScheduleStatusResult,
+        spec_set=False,
+        instance=True)
     task = cls.create_mock_task('hello', 0, 1000, scheduleStatus)
     mock_query_result.result.scheduleStatusResult.tasks = [task]
     return mock_query_result
@@ -143,14 +150,13 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('twitter.common.app.get_options', return_value=mock_options),
         patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)):
 
-      mock_api = mock_api_factory.return_value
+      mock_api = mock_api_factory('x').return_value
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
         self.assertRaises(SystemExit, kill, ['west/mchucarroll/test/hello', fp.name], mock_options)
 
-      # Now check that the right API calls got made.
-      mock_api.kill_job.call_count == 0
+      assert mock_api.kill_job.call_count == 0
 
   def test_simple_successful_killall_job(self):
     """Run a test of the "kill" command against a mocked-out API:
@@ -168,19 +174,20 @@ class TestClientKillCommand(AuroraClientCommandTest):
         self.create_mock_status_query_result(ScheduleStatus.KILLED),
     ]
     mock_scheduler_proxy.getTasksWithoutConfigs.side_effect = mock_query_results
+    mock_event = create_autospec(spec=threading.Event)
     with contextlib.nested(
         patch('time.sleep'),
         patch('apache.aurora.client.commands.core.make_client',
             return_value=mock_api),
         patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)):
+        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config),
+        patch('threading._Event.wait', return_value=mock_event)):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
         killall(['west/mchucarroll/test/hello', fp.name], mock_options)
 
-      # Now check that the right API calls got made.
       self.assert_kill_job_called(mock_api)
       mock_api.kill_job.assert_called_with(
         AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV,
@@ -216,18 +223,13 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.commands.core.make_client',
             return_value=mock_api),
         patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)) as (
-            sleep,
-            mock_make_client,
-            options,
-            mock_get_job_config):
+        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
         killall(['west/mchucarroll/test/hello', fp.name], mock_options)
 
-      # Now check that the right API calls got made.
       self.assert_kill_job_called(mock_api)
       mock_api.kill_job.assert_called_with(
         AuroraJobKey(cluster=self.TEST_CLUSTER, role=self.TEST_ROLE, env=self.TEST_ENV,
@@ -259,11 +261,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.commands.core.make_client',
             return_value=mock_api),
         patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)) as (
-            sleep,
-            mock_make_client,
-            options,
-            mock_get_job_config):
+        patch('apache.aurora.client.commands.core.get_job_config', return_value=mock_config)):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
@@ -272,31 +270,28 @@ class TestClientKillCommand(AuroraClientCommandTest):
             SystemExit, killall, ['west/mchucarroll/test/hello', fp.name], mock_options)
 
       CoreCommandHook.clear_hooks()
-      mock_api.kill_job.call_count == 0
+      assert mock_api.kill_job.call_count == 0
 
   def create_status_call_result(cls):
     """Set up the mock status call that will be used to get a task list for
     a batched kill command.
     """
-    status_response = Mock(spec=Response)
-    status_response.responseCode = ResponseCode.OK
-    status_response.messageDEPRECATED = "Ok"
-    status_response.result = Mock(spec=Result)
-    schedule_status = Mock(spec=ScheduleStatusResult)
-    status_response.result.scheduleStatusResult = schedule_status
-    mock_task_config = Mock()
-    # This should be a list of ScheduledTask's.
-    schedule_status.tasks = []
+    tasks = []
     for i in range(20):
-      task_status = Mock(spec=ScheduledTask)
-      task_status.assignedTask = Mock(spec=AssignedTask)
-      task_status.assignedTask.instanceId = i
-      task_status.assignedTask.taskId = "Task%s" % i
-      task_status.assignedTask.slaveId = "Slave%s" % i
-      task_status.slaveHost = "Slave%s" % i
-      task_status.assignedTask.task = mock_task_config
-      schedule_status.tasks.append(task_status)
-    return status_response
+      tasks.append(ScheduledTask(
+        assignedTask=AssignedTask(
+          instanceId=i,
+          taskId='Task%s' % i,
+          slaveId='Slave%s' % i,
+          slaveHost='SlaveHost%s' % i,
+          task=TaskConfig()
+        )
+      ))
+    return Response(
+        responseCode=ResponseCode.OK,
+        messageDEPRECATED='Ok',
+        result=Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
+    )
 
   def test_successful_batched_killall_job(self):
     """Run a test of the "kill" command against a mocked-out API:
@@ -439,6 +434,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
     mock_options = self.setup_mock_options()
     mock_options.batch_size = 2
     mock_options.shards = [0, 1, 2, 3]
+    mock_options.max_failures_option = 3
     mock_config = Mock()
     mock_config.hooks = []
     mock_config.raw.return_value.enable_hooks.return_value.get.return_value = False
@@ -460,7 +456,7 @@ class TestClientKillCommand(AuroraClientCommandTest):
         # We should get an exception in this case, because the one of the two calls fails.
         self.assertRaises(SystemExit, kill, ['west/mchucarroll/test/hello', fp.name], mock_options)
 
-      # killTasks should still have gotten called twice - the first error shouldn't abort
+      # killTasks should still have been called twice - the first error shouldn't abort
       # the second batch.
       assert mock_scheduler_proxy.killTasks.call_count == 2
       query = self.get_expected_task_query([2, 3])

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_listjobs.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_listjobs.py b/src/test/python/apache/aurora/client/commands/test_listjobs.py
index d9fa617..a7cb98f 100644
--- a/src/test/python/apache/aurora/client/commands/test_listjobs.py
+++ b/src/test/python/apache/aurora/client/commands/test_listjobs.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import create_autospec, Mock, patch
 
 from apache.aurora.client.commands.core import list_jobs
 
@@ -46,7 +46,7 @@ class TestListJobs(AuroraClientCommandTest):
   @classmethod
   def create_listjobs_response(cls):
     resp = cls.create_simple_success_response()
-    resp.result.getJobsResult = Mock(spec=GetJobsResult)
+    resp.result.getJobsResult = create_autospec(spec=GetJobsResult, instance=True)
     resp.result.getJobsResult.configs = set(cls.create_mock_jobs())
     return resp
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_restart.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_restart.py b/src/test/python/apache/aurora/client/commands/test_restart.py
index 6da9e73..3789f3b 100644
--- a/src/test/python/apache/aurora/client/commands/test_restart.py
+++ b/src/test/python/apache/aurora/client/commands/test_restart.py
@@ -14,8 +14,9 @@
 
 import contextlib
 import functools
+from collections import namedtuple
 
-from mock import Mock, patch
+from mock import call, create_autospec, patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.api.health_check import Retriable, StatusHealthCheck
@@ -23,38 +24,53 @@ from apache.aurora.client.commands.core import restart
 
 from .util import AuroraClientCommandTest
 
+from gen.apache.aurora.api.constants import ACTIVE_STATES
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,
     JobKey,
     PopulateJobResult,
+    Result,
     ScheduledTask,
     ScheduleStatusResult,
-    TaskConfig
+    TaskConfig,
+    TaskQuery
 )
 
 
-class TestRestartCommand(AuroraClientCommandTest):
+class FakeOptions(namedtuple('FakeOptions', ['max_total_failures',
+     'disable_all_hooks_reason',
+     'batch_size',
+     'restart_threshold',
+     'watch_secs',
+     'max_per_shard_failures',
+     'shards',
+     'health_check_interval_seconds',
+     'open_browser'])):
+
+  def __new__(cls,
+      max_total_failures=None,
+      disable_all_hooks_reason=None,
+      batch_size=None,
+      restart_threshold=None,
+      watch_secs=None,
+      max_per_shard_failures=None,
+      shards=None,
+      health_check_interval_seconds=None,
+      open_browser=None):
+    return super(FakeOptions, cls).__new__(
+        cls,
+        max_total_failures,
+        disable_all_hooks_reason,
+        batch_size,
+        restart_threshold,
+        watch_secs,
+        max_per_shard_failures,
+        shards,
+        health_check_interval_seconds,
+        open_browser)
 
-  @classmethod
-  def setup_mock_options(cls):
-    """set up to get a mock options object."""
-    mock_options = Mock(
-        spec=['json', 'bindings', 'open_browser', 'shards', 'cluster',
-              'health_check_interval_seconds', 'batch_size', 'max_per_shard_failures',
-              'max_total_failures', 'restart_threshold', 'watch_secs'])
-    mock_options.json = False
-    mock_options.bindings = {}
-    mock_options.open_browser = False
-    mock_options.shards = None
-    mock_options.cluster = None
-    mock_options.health_check_interval_seconds = 3
-    mock_options.batch_size = 5
-    mock_options.max_per_shard_failures = 0
-    mock_options.max_total_failures = 0
-    mock_options.restart_threshold = 30
-    mock_options.watch_secs = 30
-    mock_options.disable_all_hooks_reason = None
-    return mock_options
+
+class TestRestartCommand(AuroraClientCommandTest):
 
   @classmethod
   def setup_mock_scheduler_for_simple_restart(cls, api):
@@ -67,61 +83,63 @@ class TestRestartCommand(AuroraClientCommandTest):
   @classmethod
   def setup_populate_job_config(cls, api):
     populate = cls.create_simple_success_response()
-    populate.result.populateJobResult = Mock(spec=PopulateJobResult)
+    populate.result.populateJobResult = create_autospec(
+        spec=PopulateJobResult,
+        spec_set=False,
+        instance=True,
+        watch_secs=None)
     api.populateJobConfig.return_value = populate
     configs = []
     for i in range(20):
-      task_config = Mock(spec=TaskConfig)
+      task_config = create_autospec(spec=TaskConfig, instance=True)
       configs.append(task_config)
     populate.result.populateJobResult.populatedDEPRECATED = set(configs)
     return populate
 
   @classmethod
   def setup_get_tasks_status_calls(cls, scheduler):
+
+    tasks = []
+    for i in range(20):
+      tasks.append(ScheduledTask(
+        assignedTask=AssignedTask(
+          slaveHost='slave%s' % i,
+          instanceId=i,
+          taskId='task%s' % i,
+          slaveId='slave%s' % i,
+          task=TaskConfig())))
     status_response = cls.create_simple_success_response()
+    status_response.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
     scheduler.getTasksWithoutConfigs.return_value = status_response
-    schedule_status = Mock(spec=ScheduleStatusResult)
-    status_response.result.scheduleStatusResult = schedule_status
-    mock_task_config = Mock()
-    # This should be a list of ScheduledTask's.
-    schedule_status.tasks = []
-    for i in range(20):
-      task_status = Mock(spec=ScheduledTask)
-      task_status.assignedTask = Mock(spec=AssignedTask)
-      task_status.assignedTask.instanceId = i
-      task_status.assignedTask.taskId = "Task%s" % i
-      task_status.assignedTask.slaveId = "Slave%s" % i
-      task_status.slaveHost = "Slave%s" % i
-      task_status.assignedTask.task = mock_task_config
-      schedule_status.tasks.append(task_status)
 
   @classmethod
   def setup_health_checks(cls, mock_api):
-    mock_health_check = Mock(spec=StatusHealthCheck)
+    mock_health_check = create_autospec(spec=StatusHealthCheck, instance=True)
     mock_health_check.health.return_value = Retriable.alive()
     return mock_health_check
 
   def test_restart_simple(self):
-    # Test the client-side restart logic in its simplest case: everything succeeds
-    mock_options = self.setup_mock_options()
+    options = FakeOptions(
+        max_total_failures=1,
+        batch_size=5,
+        restart_threshold=10,
+        watch_secs=10)
     (mock_api, mock_scheduler_proxy) = self.create_mock_api()
     mock_health_check = self.setup_health_checks(mock_api)
     self.setup_mock_scheduler_for_simple_restart(mock_api)
     with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
+        patch('twitter.common.app.get_options', return_value=options),
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
         patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck',
             return_value=mock_health_check),
         patch('time.time', side_effect=functools.partial(self.fake_time, self)),
-        patch('threading._Event.wait')
-    ) as (options, scheduler_proxy_class, test_clusters, mock_health_check_factory,
-        time_patch, sleep_patch):
+        patch('threading._Event.wait')):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
-        restart(['west/mchucarroll/test/hello'], mock_options)
+        restart(['west/mchucarroll/test/hello'], options)
 
         # Like the update test, the exact number of calls here doesn't matter.
         # what matters is that it must have been called once before batching, plus
@@ -134,43 +152,37 @@ class TestRestartCommand(AuroraClientCommandTest):
             role=self.TEST_ROLE, name=self.TEST_JOB), [15, 16, 17, 18, 19], None)
 
   def test_restart_simple_invalid_max_failures(self):
-    # Test the client-side restart logic in its simplest case: everything succeeds
-    mock_options = self.setup_mock_options()
-    mock_options.max_total_failures = -1
-    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+    options = FakeOptions(
+        max_total_failures=None,
+        batch_size=5,
+        restart_threshold=10,
+        watch_secs=10)
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
     mock_health_check = self.setup_health_checks(mock_api)
     self.setup_mock_scheduler_for_simple_restart(mock_api)
     with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
+        patch('twitter.common.app.get_options', return_value=options),
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
         patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck',
             return_value=mock_health_check),
         patch('time.time', side_effect=functools.partial(self.fake_time, self)),
-        patch('threading._Event.wait')
-    ) as (options, scheduler_proxy_class, test_clusters, mock_health_check_factory,
-        time_patch, sleep_patch):
+        patch('threading._Event.wait')):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
-        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], mock_options)
-
-        # Like the update test, the exact number of calls here doesn't matter.
-        # what matters is that it must have been called once before batching, plus
-        # at least once per batch, and there are 4 batches.
-        assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 0
-        # called once per batch
-        assert mock_scheduler_proxy.restartShards.call_count == 0
+        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], options)
+        assert mock_scheduler_proxy.mock_calls == []
 
   def test_restart_failed_status(self):
-    mock_options = self.setup_mock_options()
-    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+    options = FakeOptions()
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
     mock_health_check = self.setup_health_checks(mock_api)
     self.setup_mock_scheduler_for_simple_restart(mock_api)
     mock_scheduler_proxy.getTasksWithoutConfigs.return_value = self.create_error_response()
     with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
+        patch('twitter.common.app.get_options', return_value=options),
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
         patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck',
@@ -183,34 +195,39 @@ class TestRestartCommand(AuroraClientCommandTest):
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
-        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], mock_options)
-        assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 1
-        assert mock_scheduler_proxy.restartShards.call_count == 0
+        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], options)
+        # TODO(wfarner): Spread this pattern further, as it flags unexpected method calls.
+        assert mock_scheduler_proxy.mock_calls == [
+          call.getTasksWithoutConfigs(
+              TaskQuery(jobKeys=[JobKey('mchucarroll', 'test', 'hello')], statuses=ACTIVE_STATES))
+        ]
 
   def test_restart_failed_restart(self):
-    # Test the client-side updater logic in its simplest case: everything succeeds, and no rolling
-    # updates.
-    mock_options = self.setup_mock_options()
-    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+    options = FakeOptions(
+        max_total_failures=1,
+        batch_size=5,
+        restart_threshold=10,
+        watch_secs=10)
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
     mock_health_check = self.setup_health_checks(mock_api)
     self.setup_mock_scheduler_for_simple_restart(mock_api)
     mock_scheduler_proxy.restartShards.return_value = self.create_error_response()
     with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
+        patch('twitter.common.app.get_options', return_value=options),
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
         patch('apache.aurora.client.api.instance_watcher.StatusHealthCheck',
             return_value=mock_health_check),
         patch('time.time', side_effect=functools.partial(self.fake_time, self)),
-        patch('threading._Event.wait')
-    ) as (options, scheduler_proxy_class, test_clusters, mock_health_check_factory,
-        time_patch, sleep_patch):
+        patch('threading._Event.wait')):
 
       with temporary_file() as fp:
         fp.write(self.get_valid_config())
         fp.flush()
-        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], mock_options)
+        self.assertRaises(SystemExit, restart, ['west/mchucarroll/test/hello'], options)
         assert mock_scheduler_proxy.getTasksWithoutConfigs.call_count == 1
         assert mock_scheduler_proxy.restartShards.call_count == 1
-        mock_scheduler_proxy.restartShards.assert_called_with(JobKey(environment=self.TEST_ENV,
-            role=self.TEST_ROLE, name=self.TEST_JOB), [0, 1, 2, 3, 4], None)
+        mock_scheduler_proxy.restartShards.assert_called_with(
+            JobKey(environment=self.TEST_ENV, role=self.TEST_ROLE, name=self.TEST_JOB),
+            [0, 1, 2, 3, 4],
+            None)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/test_run.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_run.py b/src/test/python/apache/aurora/client/commands/test_run.py
index 738a42b..528c3e9 100644
--- a/src/test/python/apache/aurora/client/commands/test_run.py
+++ b/src/test/python/apache/aurora/client/commands/test_run.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import create_autospec, Mock, patch
 
 from apache.aurora.client.commands.run import run
 
@@ -26,6 +26,7 @@ from gen.apache.aurora.api.ttypes import (
     Identity,
     JobKey,
     ResponseCode,
+    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
     TaskConfig,
@@ -49,39 +50,41 @@ class TestRunCommand(AuroraClientCommandTest):
 
   @classmethod
   def create_mock_scheduled_tasks(cls):
-    jobs = []
+    tasks = []
     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='mchucarroll')
-      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
+      task = ScheduledTask(
+        failureCount=0,
+        status=ScheduleStatus.RUNNING,
+        taskEvents=[
+            TaskEvent(timestamp=123, status=ScheduleStatus.RUNNING, message='Fake message')],
+        assignedTask=AssignedTask(
+            assignedPorts={},
+            slaveHost='slavehost',
+            instanceId=0,
+            taskId='taskid',
+            task=TaskConfig(
+                maxTaskFailures=1,
+                executorConfig='fake data',
+                metadata=[],
+                owner=Identity(role='fakerole'),
+                environment='test',
+                jobName=name,
+                numCpus=2,
+                ramMb=2,
+                diskMb=2
+            )
+        )
+      )
+      tasks.append(task)
+    return tasks
 
   @classmethod
   def create_status_response(cls):
     resp = cls.create_simple_success_response()
-    resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
+    resp.result.scheduleStatusResult = create_autospec(
+        spec=ScheduleStatusResult,
+        spec_set=False,
+        instance=True)
     resp.result.scheduleStatusResult.tasks = cls.create_mock_scheduled_tasks()
     return resp
 
@@ -127,6 +130,6 @@ class TestRunCommand(AuroraClientCommandTest):
       # The mock status call returns 3 three ScheduledTasks, so three commands should have been run
       assert mock_subprocess.call_count == 3
       mock_subprocess.assert_called_with(['ssh', '-n', '-q', 'mchucarroll@slavehost',
-          'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
+          'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-taskid/runs/'
           'slaverun/sandbox;ls'],
           stderr=-2, stdout=-1)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/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 6d51bb3..5af9133 100644
--- a/src/test/python/apache/aurora/client/commands/test_ssh.py
+++ b/src/test/python/apache/aurora/client/commands/test_ssh.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import create_autospec, Mock, patch
 
 from apache.aurora.client.commands.ssh import ssh
 
@@ -54,10 +54,10 @@ class TestSshCommand(AuroraClientCommandTest):
       job = Mock()
       job.key = job_key
       job.failure_count = 0
-      job.assignedTask = Mock(spec=AssignedTask)
+      job.assignedTask = create_autospec(spec=AssignedTask, instance=True)
       job.assignedTask.taskId = 1287391823
       job.assignedTask.slaveHost = 'slavehost'
-      job.assignedTask.task = Mock(spec=TaskConfig)
+      job.assignedTask.task = create_autospec(spec=TaskConfig, instance=True)
       job.assignedTask.task.executorConfig = Mock()
       job.assignedTask.task.maxTaskFailures = 1
       job.assignedTask.task.metadata = []
@@ -71,7 +71,7 @@ class TestSshCommand(AuroraClientCommandTest):
       job.assignedTask.instanceId = 4237894
       job.assignedTask.assignedPorts = {}
       job.status = ScheduleStatus.RUNNING
-      mockEvent = Mock(spec=TaskEvent)
+      mockEvent = create_autospec(spec=TaskEvent, instance=True)
       mockEvent.timestamp = 28234726395
       mockEvent.status = ScheduleStatus.RUNNING
       mockEvent.message = "Hi there"

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/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 062b132..003af25 100644
--- a/src/test/python/apache/aurora/client/commands/test_status.py
+++ b/src/test/python/apache/aurora/client/commands/test_status.py
@@ -14,7 +14,7 @@
 
 import contextlib
 
-from mock import Mock, patch
+from mock import patch
 
 from apache.aurora.client.commands.core import status
 
@@ -25,6 +25,8 @@ from gen.apache.aurora.api.ttypes import (
     Identity,
     JobKey,
     ResponseCode,
+    Result,
+    ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
     TaskConfig,
@@ -35,44 +37,30 @@ from gen.apache.aurora.api.ttypes import (
 
 class TestListJobs(AuroraClientCommandTest):
   @classmethod
-  def setup_mock_options(cls):
-    """set up to get a mock options object."""
-    mock_options = Mock()
-    mock_options.pretty = False
-    mock_options.show_cron = False
-    mock_options.disable_all_hooks = False
-    return mock_options
-
-  @classmethod
   def create_mock_scheduled_tasks(cls):
-    jobs = []
+    tasks = []
     for name in ['foo', 'bar', 'baz']:
-      job_key = JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV, name=name)
-      job = Mock()
-      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.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
-      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
+      tasks.append(ScheduledTask(
+          status=ScheduleStatus.RUNNING,
+          failureCount=0,
+          taskEvents=[TaskEvent(timestamp=123, status=ScheduleStatus.RUNNING, message='Hi there')],
+          assignedTask=AssignedTask(
+              instanceId=0,
+              assignedPorts={},
+              task=TaskConfig(
+                  maxTaskFailures=1,
+                  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
+              )
+          )
+      ))
+    return tasks
 
   @classmethod
   def create_mock_scheduled_task_no_metadata(cls):
@@ -84,15 +72,16 @@ class TestListJobs(AuroraClientCommandTest):
   @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 = Result(
+        scheduleStatusResult=ScheduleStatusResult(tasks=set(cls.create_mock_scheduled_tasks())))
     return resp
 
   @classmethod
   def create_status_null_metadata(cls):
     resp = cls.create_simple_success_response()
-    resp.result.scheduleStatusResult = Mock(spec=ScheduleStatusResult)
-    resp.result.scheduleStatusResult.tasks = set(cls.create_mock_scheduled_task_no_metadata())
+    resp.result = Result(
+        scheduleStatusResult=ScheduleStatusResult(
+            tasks=set(cls.create_mock_scheduled_task_no_metadata())))
     return resp
 
   @classmethod
@@ -102,15 +91,15 @@ class TestListJobs(AuroraClientCommandTest):
   def test_successful_status(self):
     """Test the status command."""
     # Calls api.check_status, which calls scheduler_proxy.getJobs
-    mock_options = self.setup_mock_options()
+    fake_options = {}
     (mock_api, mock_scheduler_proxy) = self.create_mock_api()
     mock_scheduler_proxy.getTasksWithoutConfigs.return_value = self.create_status_response()
     with contextlib.nested(
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
-        patch('twitter.common.app.get_options', return_value=mock_options)):
+        patch('twitter.common.app.get_options', return_value=fake_options)):
 
-      status(['west/mchucarroll/test/hello'], mock_options)
+      status(['west/mchucarroll/test/hello'], fake_options)
       # The status command sends a getTasksWithoutConfigs query to the scheduler,
       # and then prints the result.
       mock_scheduler_proxy.getTasksWithoutConfigs.assert_called_with(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/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 e359d32..3f5f83e 100644
--- a/src/test/python/apache/aurora/client/commands/test_update.py
+++ b/src/test/python/apache/aurora/client/commands/test_update.py
@@ -15,7 +15,7 @@
 import contextlib
 import functools
 
-from mock import Mock, patch
+from mock import create_autospec, patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.api.health_check import Retriable, StatusHealthCheck
@@ -35,8 +35,10 @@ from gen.apache.aurora.api.ttypes import (
     AssignedTask,
     JobConfiguration,
     JobKey,
+    Lock,
     PopulateJobResult,
     ResponseCode,
+    Result,
     ScheduledTask,
     ScheduleStatus,
     ScheduleStatusResult,
@@ -56,21 +58,20 @@ class TestUpdateCommand(AuroraClientCommandTest):
   @classmethod
   def setup_mock_options(cls):
     """set up to get a mock options object."""
-    mock_options = Mock()
-    mock_options.json = False
-    mock_options.bindings = {}
-    mock_options.open_browser = False
-    mock_options.cluster = None
-    mock_options.force = True
-    mock_options.env = None
-    mock_options.shards = None
-    mock_options.health_check_interval_seconds = 3
-    mock_options.disable_all_hooks = False
+    mock_options = create_autospec(
+        spec=object,
+        spec_set=False,
+        instance=True,
+        json=False,
+        shards=None,
+        health_check_interval_seconds=3,
+        force=True,
+        disable_all_hooks_reason='Fake reason')
     return mock_options
 
   @classmethod
   def setup_mock_updater(cls):
-    updater = Mock(spec=Updater)
+    updater = create_autospec(spec=Updater, instance=True)
     return updater
 
   # First, we pretend that the updater isn't really client-side, and test
@@ -117,8 +118,7 @@ class TestUpdateCommand(AuroraClientCommandTest):
     """Set up all of the API mocks for scheduler calls during a simple update"""
     sched_proxy = api.scheduler_proxy
     # First, the updater acquires a lock
-    sched_proxy.acquireLock.return_value = cls.create_acquire_lock_response(ResponseCode.OK,
-         'OK', 'token', False)
+    sched_proxy.acquireLock.return_value = cls.create_acquire_lock_response(ResponseCode.OK, 'OK')
     # Then in gets the status of the tasks for the updating job.
     cls.setup_get_tasks_status_calls(sched_proxy)
     # Next, it needs to populate the update config.
@@ -145,7 +145,7 @@ class TestUpdateCommand(AuroraClientCommandTest):
   @classmethod
   def setup_populate_job_config(cls, api):
     populate = cls.create_simple_success_response()
-    populate.result.populateJobResult = Mock(spec=PopulateJobResult)
+
     api.populateJobConfig.return_value = populate
     configs = []
     for _ in range(20):
@@ -154,18 +154,18 @@ class TestUpdateCommand(AuroraClientCommandTest):
           ramMb=1,
           diskMb=1,
           job=JobKey(role='mchucarroll', environment='test', name='hello'))
-
       configs.append(task_config)
-    populate.result.populateJobResult.populatedDEPRECATED = set(configs)
+    populate.result = Result(populateJobResult=PopulateJobResult(
+        populatedDEPRECATED=set(configs)
+    ))
     return populate
 
   @classmethod
-  def create_acquire_lock_response(cls, code, msg, token, rolling):
+  def create_acquire_lock_response(cls, code, msg):
     """Set up the response to a startUpdate API call."""
     start_update_response = cls.create_blank_response(code, msg)
-    start_update_response.result.acquireLockResult = Mock(spec=AcquireLockResult)
-    start_update_response.result.acquireLockResult.lock = "foo"
-    start_update_response.result.acquireLockResult.updateToken = 'token'
+    start_update_response.result = Result(
+        acquireLockResult=AcquireLockResult(lock=Lock(key='foo', token='token')))
     return start_update_response
 
   @classmethod
@@ -180,7 +180,7 @@ class TestUpdateCommand(AuroraClientCommandTest):
     status_response = cls.create_simple_success_response()
     scheduler_proxy.getTasksStatus.return_value = status_response
     scheduler_proxy.getTasksWithoutConfigs.return_value = status_response
-    schedule_status = Mock(spec=ScheduleStatusResult)
+    schedule_status = create_autospec(spec=ScheduleStatusResult, instance=True)
     status_response.result.scheduleStatusResult = schedule_status
     task_config = TaskConfig(
         numCpus=1.0,
@@ -191,8 +191,8 @@ class TestUpdateCommand(AuroraClientCommandTest):
     # This should be a list of ScheduledTask's.
     schedule_status.tasks = []
     for i in range(20):
-      task_status = Mock(spec=ScheduledTask)
-      task_status.assignedTask = Mock(spec=AssignedTask)
+      task_status = create_autospec(spec=ScheduledTask, instance=True)
+      task_status.assignedTask = create_autospec(spec=AssignedTask, instance=True)
       task_status.assignedTask.instanceId = i
       task_status.assignedTask.taskId = "Task%s" % i
       task_status.assignedTask.slaveId = "Slave%s" % i
@@ -207,20 +207,20 @@ class TestUpdateCommand(AuroraClientCommandTest):
 
   @classmethod
   def setup_health_checks(cls, mock_api):
-    mock_health_check = Mock(spec=StatusHealthCheck)
+    mock_health_check = create_autospec(spec=StatusHealthCheck, instance=True)
     mock_health_check.health.return_value = Retriable.alive()
     return mock_health_check
 
   @classmethod
   def setup_quota_check(cls):
-    mock_quota_check = Mock(spec=QuotaCheck)
+    mock_quota_check = create_autospec(spec=QuotaCheck, instance=True)
     mock_quota_check.validate_quota_from_requested.return_value = (
         cls.create_simple_success_response())
     return mock_quota_check
 
   @classmethod
   def setup_job_monitor(cls):
-    mock_job_monitor = Mock(spec=JobMonitor)
+    mock_job_monitor = create_autospec(spec=JobMonitor, instance=True)
     mock_job_monitor.wait_until.return_value = True
     return mock_job_monitor
 
@@ -286,7 +286,7 @@ class TestUpdateCommand(AuroraClientCommandTest):
             jobKeys=[JobKey(role='mchucarroll', environment='test', name='hello')],
             instanceIds=frozenset([19]),
             statuses=ACTIVE_STATES),
-        'foo')
+        Lock(key='foo', token='token'))
 
   @classmethod
   def assert_correct_status_calls(cls, api):
@@ -302,7 +302,7 @@ class TestUpdateCommand(AuroraClientCommandTest):
       status_call[0][0] == TaskQuery(
         taskIds=None,
         jobKeys=[JobKey(role='mchucarroll', environment='test', name='hello')],
-        statuses=set([ScheduleStatus.RUNNING]))
+        statuses={ScheduleStatus.RUNNING})
 
     # getTasksStatus is called only once to build an generate update instructions
     assert api.getTasksStatus.call_count == 1

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/commands/util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/util.py b/src/test/python/apache/aurora/client/commands/util.py
index fb84258..c06de50 100644
--- a/src/test/python/apache/aurora/client/commands/util.py
+++ b/src/test/python/apache/aurora/client/commands/util.py
@@ -15,7 +15,7 @@
 import unittest
 from collections import defaultdict
 
-from mock import Mock
+from mock import create_autospec
 
 from apache.aurora.client.api.sla import DomainUpTimeSlaVector, JobUpTimeDetails
 from apache.aurora.client.hooks.hooked_api import HookedAuroraClientAPI
@@ -23,16 +23,19 @@ from apache.aurora.common.aurora_job_key import AuroraJobKey
 from apache.aurora.common.cluster import Cluster
 from apache.aurora.common.clusters import Clusters
 
+from ..api.api_util import SchedulerProxyApiSpec
+
 from gen.apache.aurora.api.ttypes import Response, ResponseCode, Result
 
 
 class AuroraClientCommandTest(unittest.TestCase):
   @classmethod
   def create_blank_response(cls, code, msg):
-    response = Mock(spec=Response)
+    # TODO(wfarner): Don't use a mock here.
+    response = create_autospec(spec=Response, instance=True)
     response.responseCode = code
     response.messageDEPRECATED = msg
-    response.result = Mock(spec=Result)
+    response.result = create_autospec(spec=Result, instance=True)
     response.details = []
     return response
 
@@ -47,25 +50,24 @@ class AuroraClientCommandTest(unittest.TestCase):
   @classmethod
   def create_mock_api(cls):
     """Builds up a mock API object, with a mock SchedulerProxy"""
-    # This looks strange, but we set up the same object to use as both
-    # the SchedulerProxy and the SchedulerClient. These tests want to observe
-    # what API calls get made against the scheduler, and both of these objects
-    # delegate calls to the scheduler. It doesn't matter which one is used:
-    # what we care about is that the right API calls get made.
-    mock_scheduler_proxy = Mock()
+    """Builds up a mock API object, with a mock SchedulerProxy.
+    Returns the API and the proxy"""
+    mock_scheduler_proxy = create_autospec(
+        spec=SchedulerProxyApiSpec,
+        spec_set=False,
+        instance=True)
     mock_scheduler_proxy.url = "http://something_or_other"
     mock_scheduler_proxy.scheduler_client.return_value = mock_scheduler_proxy
-    mock_api = Mock(spec=HookedAuroraClientAPI)
+    mock_api = create_autospec(spec=HookedAuroraClientAPI)
     mock_api.scheduler_proxy = mock_scheduler_proxy
-    return (mock_api, mock_scheduler_proxy)
+    return mock_api, mock_scheduler_proxy
 
   @classmethod
   def create_mock_api_factory(cls):
     """Create a collection of mocks for a test that wants to mock out the client API
     by patching the api factory."""
     mock_api, mock_scheduler_proxy = cls.create_mock_api()
-    mock_api_factory = Mock()
-    mock_api_factory.return_value = mock_api
+    mock_api_factory = lambda x: mock_api
     return mock_api_factory, mock_scheduler_proxy
 
   FAKE_TIME = 42131
@@ -132,7 +134,7 @@ jobs = [HELLO_WORLD]
 
   @classmethod
   def create_mock_probe_hosts_vector(cls, side_effects):
-    mock_vector = Mock(spec=DomainUpTimeSlaVector)
+    mock_vector = create_autospec(spec=DomainUpTimeSlaVector, instance=True)
     mock_vector.probe_hosts.side_effect = side_effects
     return mock_vector
 
@@ -143,8 +145,9 @@ jobs = [HELLO_WORLD]
     hosts[hostname].append(JobUpTimeDetails(job, predicted, safe, safe_in))
     return [hosts]
 
+  #TODO(wfarner): Remove this, force tests to call out their flags.
   @classmethod
   def setup_mock_options(cls):
-    mock_options = Mock(spec=['verbosity'])
+    mock_options = create_autospec(spec=['verbosity'], instance=True)
     mock_options.verbosity = 'verbose'
     return mock_options

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/client/hooks/test_hooked_api.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/hooks/test_hooked_api.py b/src/test/python/apache/aurora/client/hooks/test_hooked_api.py
index 496bdb0..c2063f2 100644
--- a/src/test/python/apache/aurora/client/hooks/test_hooked_api.py
+++ b/src/test/python/apache/aurora/client/hooks/test_hooked_api.py
@@ -14,7 +14,7 @@
 
 from inspect import getargspec
 
-from mock import Mock
+from mock import create_autospec, Mock
 
 from apache.aurora.client.api import AuroraClientAPI
 from apache.aurora.client.hooks.hooked_api import HookedAuroraClientAPI, NonHookedAuroraClientAPI
@@ -37,7 +37,7 @@ def test_api_methods_exist(method_name):
 
 
 def test_api_methods_params(method_name):
-  cluster = Mock(spec=Cluster)
+  cluster = create_autospec(spec=Cluster, instance=True)
   api = HookedAuroraClientAPI(cluster=cluster)  # cant use mock here; need to inspect methods
 
   hooked_method = getattr(api, method_name)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/common/test_transport.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/common/test_transport.py b/src/test/python/apache/aurora/common/test_transport.py
index e0fef91..dd066bf 100644
--- a/src/test/python/apache/aurora/common/test_transport.py
+++ b/src/test/python/apache/aurora/common/test_transport.py
@@ -15,8 +15,8 @@
 import logging
 from threading import Thread
 
-import mock
 import requests
+from mock import create_autospec, Mock
 from requests import exceptions as request_exceptions
 from thrift.protocol import TJSONProtocol
 from thrift.server import THttpServer
@@ -43,8 +43,6 @@ def test_request_transport_integration():
   server_thread.start()
   _, server_port = server.httpd.socket.getsockname()
 
-  response = None
-
   try:
     transport = TRequestsTransport('http://localhost:%d' % server_port)
     protocol = TJSONProtocol.TJSONProtocol(transport)
@@ -62,9 +60,9 @@ def test_request_transport_integration():
 
 
 def test_request_transport_timeout():
-  session = mock.MagicMock(spec=requests.Session)
+  session = create_autospec(spec=requests.Session, instance=True)
   session.headers = {}
-  session.post = mock.Mock(side_effect=request_exceptions.Timeout())
+  session.post = Mock(side_effect=request_exceptions.Timeout())
   transport = TRequestsTransport('http://localhost:12345', session_factory=lambda: session)
   protocol = TJSONProtocol.TJSONProtocol(transport)
   client = ReadOnlyScheduler.Client(protocol)
@@ -81,9 +79,9 @@ def test_request_transport_timeout():
 
 
 def test_request_any_other_exception():
-  session = mock.MagicMock(spec=requests.Session)
+  session = create_autospec(spec=requests.Session, instance=True)
   session.headers = {}
-  session.post = mock.Mock(side_effect=request_exceptions.ConnectionError())
+  session.post = Mock(side_effect=request_exceptions.ConnectionError())
   transport = TRequestsTransport('http://localhost:12345', session_factory=lambda: session)
   protocol = TJSONProtocol.TJSONProtocol(transport)
   client = ReadOnlyScheduler.Client(protocol)
@@ -104,6 +102,6 @@ def test_requests_transports_lowers_logging_level():
 
   TRequestsTransport(
       'http://localhost:12345',
-      session_factory=lambda: mock.MagicMock(spec=requests.Session))
+      session_factory=lambda: create_autospec(spec=requests.Session, instance=True))
 
   assert logging.getLogger('requests').level == logging.WARNING

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/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 a4ab532..e329a90 100644
--- a/src/test/python/apache/aurora/executor/common/test_announcer.py
+++ b/src/test/python/apache/aurora/executor/common/test_announcer.py
@@ -14,10 +14,10 @@
 
 import threading
 
-import mock
 import pytest
 from kazoo.client import KazooClient
 from kazoo.exceptions import KazooException
+from mock import create_autospec, MagicMock, patch
 from twitter.common.quantity import Amount, Time
 from twitter.common.testing.clock import ThreadedClock
 from twitter.common.zookeeper.serverset import Endpoint, ServerSet
@@ -72,10 +72,10 @@ def test_announcer_under_normal_circumstances():
     joined.set()
     return 'membership foo'
 
-  mock_serverset = mock.MagicMock(spec=ServerSet)
-  mock_serverset.join = mock.MagicMock()
+  mock_serverset = create_autospec(spec=ServerSet, instance=True)
+  mock_serverset.join = MagicMock()
   mock_serverset.join.side_effect = joined_side_effect
-  mock_serverset.cancel = mock.MagicMock()
+  mock_serverset.cancel = MagicMock()
 
   endpoint = Endpoint('localhost', 12345)
   clock = ThreadedClock(31337.0)
@@ -125,10 +125,10 @@ def test_announcer_on_expiration():
     else:
       raise KazooException('Failed to reconnect')
 
-  mock_serverset = mock.MagicMock(spec=ServerSet)
-  mock_serverset.join = mock.MagicMock()
+  mock_serverset = create_autospec(spec=ServerSet, instance=True)
+  mock_serverset.join = MagicMock()
   mock_serverset.join.side_effect = joined_side_effect
-  mock_serverset.cancel = mock.MagicMock()
+  mock_serverset.cancel = MagicMock()
 
   endpoint = Endpoint('localhost', 12345)
   clock = ThreadedClock(31337.0)
@@ -160,13 +160,13 @@ def test_announcer_on_expiration():
 # TODO(wickman) https://issues.apache.org/jira/browse/AURORA-639
 @pytest.mark.skipif('True')
 def test_announcer_under_abnormal_circumstances():
-  mock_serverset = mock.MagicMock(spec=ServerSet)
-  mock_serverset.join = mock.MagicMock()
+  mock_serverset = create_autospec(spec=ServerSet, instance=True)
+  mock_serverset.join = MagicMock()
   mock_serverset.join.side_effect = [
       KazooException('Whoops the ensemble is down!'),
       'member0001',
   ]
-  mock_serverset.cancel = mock.MagicMock()
+  mock_serverset.cancel = MagicMock()
 
   endpoint = Endpoint('localhost', 12345)
   clock = ThreadedClock(31337.0)
@@ -244,15 +244,15 @@ def test_make_empty_endpoints():
   assert additional == {}
 
 
-@mock.patch('apache.aurora.executor.common.announcer.ServerSet')
-@mock.patch('apache.aurora.executor.common.announcer.KazooClient')
+@patch('apache.aurora.executor.common.announcer.ServerSet')
+@patch('apache.aurora.executor.common.announcer.KazooClient')
 def test_announcer_provider_with_timeout(mock_client_provider, mock_serverset_provider):
-  mock_client = mock.MagicMock(spec=KazooClient)
+  mock_client = create_autospec(spec=KazooClient, instance=True)
   mock_client_provider.return_value = mock_client
   client_connect_event = threading.Event()
   mock_client.start_async.return_value = client_connect_event
 
-  mock_serverset = mock.MagicMock(spec=ServerSet)
+  mock_serverset = create_autospec(spec=ServerSet, instance=True)
   mock_serverset_provider.return_value = mock_serverset
 
   dap = DefaultAnnouncerCheckerProvider('zookeeper.example.com', root='/aurora')
@@ -272,12 +272,12 @@ def test_announcer_provider_with_timeout(mock_client_provider, mock_serverset_pr
   assert checker.status is not None
 
 
-@mock.patch('apache.aurora.executor.common.announcer.ServerSet')
-@mock.patch('apache.aurora.executor.common.announcer.KazooClient')
+@patch('apache.aurora.executor.common.announcer.ServerSet')
+@patch('apache.aurora.executor.common.announcer.KazooClient')
 def test_default_announcer_provider(mock_client_provider, mock_serverset_provider):
-  mock_client = mock.MagicMock(spec=KazooClient)
+  mock_client = create_autospec(spec=KazooClient, instance=True)
   mock_client_provider.return_value = mock_client
-  mock_serverset = mock.MagicMock(spec=ServerSet)
+  mock_serverset = create_autospec(spec=ServerSet, instance=True)
   mock_serverset_provider.return_value = mock_serverset
 
   dap = DefaultAnnouncerCheckerProvider('zookeeper.example.com', root='/aurora')

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/executor/common/test_executor_timeout.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/executor/common/test_executor_timeout.py b/src/test/python/apache/aurora/executor/common/test_executor_timeout.py
index ce4cfc8..275981e 100644
--- a/src/test/python/apache/aurora/executor/common/test_executor_timeout.py
+++ b/src/test/python/apache/aurora/executor/common/test_executor_timeout.py
@@ -25,7 +25,7 @@ from apache.aurora.executor.common.executor_timeout import ExecutorTimeout
 class TestExecutorTimeout(TestCase):
   def test_run(self):
     event = Event()
-    mock_driver = mock.Mock(spec=ExecutorDriver)
+    mock_driver = mock.create_autospec(spec=ExecutorDriver, instance=True)
     event.set()
     executor_timeout = ExecutorTimeout(event, mock_driver, timeout=Amount(0, Time.SECONDS))
     executor_timeout.run()
@@ -33,7 +33,7 @@ class TestExecutorTimeout(TestCase):
 
   def test_run_timeout(self):
     event = Event()
-    mock_driver = mock.Mock(spec=ExecutorDriver)
+    mock_driver = mock.create_autospec(spec=ExecutorDriver, instance=True)
     executor_timeout = ExecutorTimeout(event, mock_driver, timeout=Amount(0, Time.SECONDS))
     executor_timeout.run()
     mock_driver.stop.assert_called_once_with()

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/aurora/executor/test_status_manager.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/executor/test_status_manager.py b/src/test/python/apache/aurora/executor/test_status_manager.py
index 8a0d48b..ce4679b 100644
--- a/src/test/python/apache/aurora/executor/test_status_manager.py
+++ b/src/test/python/apache/aurora/executor/test_status_manager.py
@@ -43,7 +43,7 @@ class TestStatusManager(TestCase):
     def callback(result):
       assert result == TaskState.Value('TASK_KILLED')
       self.callback_called = True
-    mock_time = mock.Mock(spec=time)
+    mock_time = mock.create_autospec(spec=time, instance=True)
     status_manager = StatusManager(checker, callback, mock_time)
     status_manager.run()
     assert mock_time.sleep.call_count == 2

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/d01f0c38/src/test/python/apache/thermos/core/test_helper.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/thermos/core/test_helper.py b/src/test/python/apache/thermos/core/test_helper.py
index fe368ff..35397ab 100644
--- a/src/test/python/apache/thermos/core/test_helper.py
+++ b/src/test/python/apache/thermos/core/test_helper.py
@@ -38,7 +38,7 @@ def set_side_effect(mock_obj, value):
 
 
 def mock_process(pid, username, uid=None):
-  process = mock.Mock(spec=psutil.Process, pid=pid)
+  process = mock.create_autospec(spec=psutil.Process, pid=pid, instance=True)
   set_side_effect(process.uids, uid)
   set_side_effect(process.username, username)
   process.create_time.return_value = CREATE_TIME


Mime
View raw message