Repository: aurora
Updated Branches:
  refs/heads/master 8524dbfef -> 7b1c2b2fc


http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
----------------------------------------------------------------------
diff --git 
a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java 
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
index a6769e6..2de1783 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
@@ -29,7 +29,6 @@ import org.apache.aurora.gen.ResourceAggregate;
 import org.apache.aurora.gen.Response;
 import org.apache.aurora.gen.RewriteConfigsRequest;
 import org.apache.aurora.gen.ScheduleStatus;
-import org.apache.aurora.gen.SessionKey;
 import org.apache.aurora.gen.TaskQuery;
 import org.apache.thrift.TException;
 
@@ -50,60 +49,49 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   @Override
   public Response setQuota(
       String ownerRole,
-      ResourceAggregate resourceAggregate,
-      SessionKey session) throws TException {
+      ResourceAggregate resourceAggregate) throws TException {
 
-    return delegate.setQuota(ownerRole, resourceAggregate, session);
+    return delegate.setQuota(ownerRole, resourceAggregate);
   }
 
   @Override
-  public Response forceTaskState(
-      String taskId,
-      ScheduleStatus status,
-      SessionKey session) throws TException {
-
-    return delegate.forceTaskState(taskId, status, session);
+  public Response forceTaskState(String taskId, ScheduleStatus status) throws 
TException {
+    return delegate.forceTaskState(taskId, status);
   }
 
   @Override
-  public Response performBackup(SessionKey session) throws TException {
-    return delegate.performBackup(session);
+  public Response performBackup() throws TException {
+    return delegate.performBackup();
   }
 
   @Override
-  public Response listBackups(SessionKey session) throws TException {
-    return delegate.listBackups(session);
+  public Response listBackups() throws TException {
+    return delegate.listBackups();
   }
 
   @Override
-  public Response stageRecovery(String backupId, SessionKey session)
-      throws TException {
-
-    return delegate.stageRecovery(backupId, session);
+  public Response stageRecovery(String backupId) throws TException {
+    return delegate.stageRecovery(backupId);
   }
 
   @Override
-  public Response queryRecovery(TaskQuery query, SessionKey session)
-      throws TException {
-
-    return delegate.queryRecovery(query, session);
+  public Response queryRecovery(TaskQuery query) throws TException {
+    return delegate.queryRecovery(query);
   }
 
   @Override
-  public Response deleteRecoveryTasks(TaskQuery query, SessionKey session)
-      throws TException {
-
-    return delegate.deleteRecoveryTasks(query, session);
+  public Response deleteRecoveryTasks(TaskQuery query) throws TException {
+    return delegate.deleteRecoveryTasks(query);
   }
 
   @Override
-  public Response commitRecovery(SessionKey session) throws TException {
-    return delegate.commitRecovery(session);
+  public Response commitRecovery() throws TException {
+    return delegate.commitRecovery();
   }
 
   @Override
-  public Response unloadRecovery(SessionKey session) throws TException {
-    return delegate.unloadRecovery(session);
+  public Response unloadRecovery() throws TException {
+    return delegate.unloadRecovery();
   }
 
   @Override
@@ -122,29 +110,23 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response createJob(JobConfiguration description, Lock lock, 
SessionKey session)
-      throws TException {
-
-    return delegate.createJob(description, lock, session);
+  public Response createJob(JobConfiguration description, Lock lock) throws 
TException {
+    return delegate.createJob(description, lock);
   }
 
   @Override
-  public Response scheduleCronJob(JobConfiguration description, Lock lock, 
SessionKey session)
-      throws TException {
-
-    return delegate.scheduleCronJob(description, lock, session);
+  public Response scheduleCronJob(JobConfiguration description, Lock lock) 
throws TException {
+    return delegate.scheduleCronJob(description, lock);
   }
 
   @Override
-  public Response descheduleCronJob(JobKey job, Lock lock, SessionKey session) 
throws TException {
-    return delegate.descheduleCronJob(job, lock, session);
+  public Response descheduleCronJob(JobKey job, Lock lock) throws TException {
+    return delegate.descheduleCronJob(job, lock);
   }
 
   @Override
-  public Response replaceCronTemplate(JobConfiguration config, Lock lock, 
SessionKey session)
-      throws TException {
-
-    return delegate.replaceCronTemplate(config, lock, session);
+  public Response replaceCronTemplate(JobConfiguration config, Lock lock) 
throws TException {
+    return delegate.replaceCronTemplate(config, lock);
   }
 
   @Override
@@ -153,18 +135,13 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response startCronJob(JobKey job, SessionKey session) throws 
TException {
-    return delegate.startCronJob(job, session);
+  public Response startCronJob(JobKey job) throws TException {
+    return delegate.startCronJob(job);
   }
 
   @Override
-  public Response restartShards(
-      JobKey job,
-      Set<Integer> shardIds,
-      Lock lock,
-      SessionKey session) throws TException {
-
-    return delegate.restartShards(job, shardIds, lock, session);
+  public Response restartShards(JobKey job, Set<Integer> shardIds, Lock lock) 
throws TException {
+    return delegate.restartShards(job, shardIds, lock);
   }
 
   @Override
@@ -183,8 +160,8 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response killTasks(TaskQuery query, Lock lock, SessionKey session) 
throws TException {
-    return delegate.killTasks(query, lock, session);
+  public Response killTasks(TaskQuery query, Lock lock) throws TException {
+    return delegate.killTasks(query, lock);
   }
 
   @Override
@@ -193,51 +170,43 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response startMaintenance(Hosts hosts, SessionKey session)
-      throws TException {
-
-    return delegate.startMaintenance(hosts, session);
+  public Response startMaintenance(Hosts hosts) throws TException {
+    return delegate.startMaintenance(hosts);
   }
 
   @Override
-  public Response drainHosts(Hosts hosts, SessionKey session) throws 
TException {
-    return delegate.drainHosts(hosts, session);
+  public Response drainHosts(Hosts hosts) throws TException {
+    return delegate.drainHosts(hosts);
   }
 
   @Override
-  public Response maintenanceStatus(Hosts hosts, SessionKey session)
-      throws TException {
-
-    return delegate.maintenanceStatus(hosts, session);
+  public Response maintenanceStatus(Hosts hosts) throws TException {
+    return delegate.maintenanceStatus(hosts);
   }
 
   @Override
-  public Response endMaintenance(Hosts hosts, SessionKey session) throws 
TException {
-    return delegate.endMaintenance(hosts, session);
+  public Response endMaintenance(Hosts hosts) throws TException {
+    return delegate.endMaintenance(hosts);
   }
 
   @Override
-  public Response snapshot(SessionKey session) throws TException {
-    return delegate.snapshot(session);
+  public Response snapshot() throws TException {
+    return delegate.snapshot();
   }
 
   @Override
-  public Response rewriteConfigs(RewriteConfigsRequest request, SessionKey 
session)
-      throws TException {
-
-    return delegate.rewriteConfigs(request, session);
+  public Response rewriteConfigs(RewriteConfigsRequest request) throws 
TException {
+    return delegate.rewriteConfigs(request);
   }
 
   @Override
-  public Response acquireLock(LockKey lockKey, SessionKey session) throws 
TException {
-    return delegate.acquireLock(lockKey, session);
+  public Response acquireLock(LockKey lockKey) throws TException {
+    return delegate.acquireLock(lockKey);
   }
 
   @Override
-  public Response releaseLock(Lock lock, LockValidation validation, SessionKey 
session)
-      throws TException {
-
-    return delegate.releaseLock(lock, validation, session);
+  public Response releaseLock(Lock lock, LockValidation validation) throws 
TException {
+    return delegate.releaseLock(lock, validation);
   }
 
   @Override
@@ -246,12 +215,8 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response addInstances(
-      AddInstancesConfig config,
-      Lock lock,
-      SessionKey session) throws TException {
-
-    return delegate.addInstances(config, lock, session);
+  public Response addInstances(AddInstancesConfig config, Lock lock) throws 
TException {
+    return delegate.addInstances(config, lock);
   }
 
   @Override
@@ -260,36 +225,28 @@ abstract class ForwardingThrift implements 
AnnotatedAuroraAdmin {
   }
 
   @Override
-  public Response startJobUpdate(JobUpdateRequest request, String message, 
SessionKey session)
-      throws TException {
-
-    return delegate.startJobUpdate(request, message, session);
+  public Response startJobUpdate(JobUpdateRequest request, String message) 
throws TException {
+    return delegate.startJobUpdate(request, message);
   }
 
   @Override
-  public Response pauseJobUpdate(JobUpdateKey key, String message, SessionKey 
session)
-      throws TException {
-
-    return delegate.pauseJobUpdate(key, message, session);
+  public Response pauseJobUpdate(JobUpdateKey key, String message) throws 
TException {
+    return delegate.pauseJobUpdate(key, message);
   }
 
   @Override
-  public Response resumeJobUpdate(JobUpdateKey key, String message, SessionKey 
session)
-      throws TException {
-
-    return delegate.resumeJobUpdate(key, message, session);
+  public Response resumeJobUpdate(JobUpdateKey key, String message) throws 
TException {
+    return delegate.resumeJobUpdate(key, message);
   }
 
   @Override
-  public Response abortJobUpdate(JobUpdateKey key, String message, SessionKey 
session)
-      throws TException {
-
-    return delegate.abortJobUpdate(key, message, session);
+  public Response abortJobUpdate(JobUpdateKey key, String message) throws 
TException {
+    return delegate.abortJobUpdate(key, message);
   }
 
   @Override
-  public Response pulseJobUpdate(JobUpdateKey key, SessionKey session) throws 
TException {
-    return delegate.pulseJobUpdate(key, session);
+  public Response pulseJobUpdate(JobUpdateKey key) throws TException {
+    return delegate.pulseJobUpdate(key);
   }
 
   @Override

http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/java/org/apache/aurora/scheduler/thrift/aop/LoggingInterceptorTest.java
----------------------------------------------------------------------
diff --git 
a/src/test/java/org/apache/aurora/scheduler/thrift/aop/LoggingInterceptorTest.java
 
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/LoggingInterceptorTest.java
index f0caa8b..8dff558 100644
--- 
a/src/test/java/org/apache/aurora/scheduler/thrift/aop/LoggingInterceptorTest.java
+++ 
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/LoggingInterceptorTest.java
@@ -22,7 +22,6 @@ import org.apache.aurora.common.testing.easymock.EasyMockTest;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.Response;
 import org.apache.aurora.gen.ResponseCode;
-import org.apache.aurora.gen.SessionKey;
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.scheduler.storage.Storage.TransientStorageException;
 import org.junit.Before;
@@ -53,10 +52,6 @@ public class LoggingInterceptorTest extends EasyMockTest {
     Response respond(JobConfiguration jobConfiguration) {
       throw new UnsupportedOperationException();
     }
-
-    Response respond(SessionKey sessionKey) {
-      throw new UnsupportedOperationException();
-    }
   }
 
   @Before

http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/python/apache/aurora/client/api/test_restarter.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_restarter.py 
b/src/test/python/apache/aurora/client/api/test_restarter.py
index f31d60d..7557144 100644
--- a/src/test/python/apache/aurora/client/api/test_restarter.py
+++ b/src/test/python/apache/aurora/client/api/test_restarter.py
@@ -34,11 +34,9 @@ from gen.apache.aurora.api.ttypes import (
     ScheduleStatus,
     ScheduleStatusResult,
     ServerInfo,
-    SessionKey,
     TaskConfig
 )
 
-SESSION_KEY = SessionKey()
 CLUSTER = 'east'
 JOB = AuroraJobKey(CLUSTER, 'johndoe', 'test', 'test_job')
 HEALTH_CHECK_INTERVAL_SECONDS = 5
@@ -73,18 +71,14 @@ class TestRestarter(MoxTestBase):
         JOB,
         UPDATER_CONFIG,
         HEALTH_CHECK_INTERVAL_SECONDS,
-        FakeSchedulerProxy(
-            Cluster(name=CLUSTER),
-            self.mock_scheduler,
-            SESSION_KEY),
+        FakeSchedulerProxy(Cluster(name=CLUSTER), self.mock_scheduler),
         self.mock_instance_watcher)
 
   def mock_restart_instances(self, instances, lock=None):
     self.mock_scheduler.restartShards(
         JOB.to_thrift(),
         instances,
-        lock,
-        SESSION_KEY).AndReturn(make_response())
+        lock).AndReturn(make_response())
     self.mock_instance_watcher.watch(instances).AndReturn([])
 
   def test_restart_one_iteration(self):
@@ -142,8 +136,7 @@ class TestRestarter(MoxTestBase):
     self.mock_scheduler.restartShards(
         JOB.to_thrift(),
         IgnoreArg(),
-        self.lock,
-        SESSION_KEY).AndReturn(response)
+        self.lock).AndReturn(response)
 
   def test_restart_instance_fails(self):
     self.mock_status_active_tasks([0, 1])
@@ -157,8 +150,7 @@ class TestRestarter(MoxTestBase):
     self.mock_scheduler.restartShards(
         JOB.to_thrift(),
         instances,
-        self.lock,
-        SESSION_KEY).AndReturn(make_response())
+        self.lock).AndReturn(make_response())
     self.mock_instance_watcher.watch(instances).AndReturn(instances)
 
   def test_restart_instances_watch_fails(self):

http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/python/apache/aurora/client/api/test_scheduler_client.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_scheduler_client.py 
b/src/test/python/apache/aurora/client/api/test_scheduler_client.py
index d244b3a..2daa96c 100644
--- a/src/test/python/apache/aurora/client/api/test_scheduler_client.py
+++ b/src/test/python/apache/aurora/client/api/test_scheduler_client.py
@@ -46,7 +46,6 @@ from gen.apache.aurora.api.ttypes import (
     ResponseDetail,
     RewriteConfigsRequest,
     ScheduleStatus,
-    SessionKey,
     TaskQuery
 )
 
@@ -66,16 +65,6 @@ def test_coverage():
           'No test defined for RPC %s' % rpc_name)
 
 
-SESSION = SessionKey()
-
-
-class TestSchedulerProxy(scheduler_client.SchedulerProxy):
-  """In testing we shouldn't use the real SSHAgentAuthenticator."""
-
-  def session_key(self):
-    return SESSION
-
-
 class TestSchedulerProxyInjection(unittest.TestCase):
   def setUp(self):
     self.mox = Mox()
@@ -95,39 +84,35 @@ class TestSchedulerProxyInjection(unittest.TestCase):
     self.mox.VerifyAll()
 
   def make_scheduler_proxy(self):
-    return TestSchedulerProxy(Cluster(name='local'))
+    return scheduler_client.SchedulerProxy(Cluster(name='local'))
 
   def test_startCronJob(self):
-    self.mock_thrift_client.startCronJob(IsA(JobKey), 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.startCronJob(IsA(JobKey)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().startCronJob(JOB_KEY)
 
   def test_createJob(self):
     self.mock_thrift_client.createJob(
-        IsA(JobConfiguration),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().createJob(JobConfiguration())
 
   def test_replaceCronTemplate(self):
     self.mock_thrift_client.replaceCronTemplate(
         IsA(JobConfiguration),
-        IsA(Lock),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(Lock)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().replaceCronTemplate(JobConfiguration(), Lock())
 
   def test_scheduleCronJob(self):
     self.mock_thrift_client.scheduleCronJob(
-        IsA(JobConfiguration),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().scheduleCronJob(JobConfiguration())
 
   def test_descheduleCronJob(self):
     self.mock_thrift_client.descheduleCronJob(
-        IsA(JobKey),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(JobKey)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().descheduleCronJob(JOB_KEY)
 
@@ -139,8 +124,7 @@ class TestSchedulerProxyInjection(unittest.TestCase):
   def test_restartShards(self):
     self.mock_thrift_client.restartShards(
         IsA(JobKey),
-        IgnoreArg(),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().restartShards(JOB_KEY, set([0]))
 
@@ -155,7 +139,7 @@ class TestSchedulerProxyInjection(unittest.TestCase):
     self.make_scheduler_proxy().getJobs(ROLE)
 
   def test_killTasks(self):
-    self.mock_thrift_client.killTasks(IsA(TaskQuery), 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.killTasks(IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().killTasks(TaskQuery())
 
@@ -168,21 +152,19 @@ class TestSchedulerProxyInjection(unittest.TestCase):
     self.mock_thrift_client.addInstances(
       IsA(JobKey),
       IgnoreArg(),
-      IsA(Lock),
-      SESSION).AndReturn(DEFAULT_RESPONSE)
+      IsA(Lock)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().addInstances(JobKey(), {}, Lock())
 
   def test_acquireLock(self):
-    self.mock_thrift_client.acquireLock(IsA(Lock), 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.acquireLock(IsA(Lock)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().acquireLock(Lock())
 
   def test_releaseLock(self):
     self.mock_thrift_client.releaseLock(
         IsA(Lock),
-        IsA(LockValidation),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(LockValidation)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().releaseLock(Lock(), LockValidation())
 
@@ -198,35 +180,33 @@ class TestSchedulerProxyInjection(unittest.TestCase):
 
   def test_startJobUpdate(self):
     self.mock_thrift_client.startJobUpdate(
-        IsA(JobUpdateRequest),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(JobUpdateRequest)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().startJobUpdate(JobUpdateRequest())
 
   def test_pauseJobUpdate(self):
-    self.mock_thrift_client.pauseJobUpdate('update_id', 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.pauseJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().pauseJobUpdate('update_id')
 
   def test_resumeJobUpdate(self):
     self.mock_thrift_client.resumeJobUpdate(
-        'update_id',
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        'update_id').AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().resumeJobUpdate('update_id')
 
   def test_abortJobUpdate(self):
-    self.mock_thrift_client.abortJobUpdate('update_id', 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.abortJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().abortJobUpdate('update_id')
 
   def test_pulseJobUpdate(self):
-    self.mock_thrift_client.pulseJobUpdate('update_id', 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.pulseJobUpdate('update_id').AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().pulseJobUpdate('update_id')
 
   def test_raise_auth_error(self):
-    self.mock_thrift_client.killTasks(TaskQuery(), None, None, 
SESSION).AndRaise(
+    self.mock_thrift_client.killTasks(TaskQuery(), None, None).AndRaise(
         TRequestsTransport.AuthError())
     self.mock_scheduler_client.get_failed_auth_message().AndReturn('failed 
auth')
     self.mox.ReplayAll()
@@ -237,94 +217,86 @@ class TestSchedulerProxyInjection(unittest.TestCase):
 class TestSchedulerProxyAdminInjection(TestSchedulerProxyInjection):
   def test_startMaintenance(self):
     self.mock_thrift_client.startMaintenance(
-      IsA(Hosts),
-      SESSION).AndReturn(DEFAULT_RESPONSE)
+      IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().startMaintenance(Hosts())
 
   def test_drainHosts(self):
-    self.mock_thrift_client.drainHosts(IsA(Hosts), 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.drainHosts(IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().drainHosts(Hosts())
 
   def test_maintenanceStatus(self):
     self.mock_thrift_client.maintenanceStatus(
-      IsA(Hosts),
-      SESSION).AndReturn(DEFAULT_RESPONSE)
+      IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().maintenanceStatus(Hosts())
 
   def test_endMaintenance(self):
-    self.mock_thrift_client.endMaintenance(IsA(Hosts), 
SESSION).AndReturn(DEFAULT_RESPONSE)
+    
self.mock_thrift_client.endMaintenance(IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().endMaintenance(Hosts())
 
   def test_setQuota(self):
     self.mock_thrift_client.setQuota(
         IgnoreArg(),
-        IsA(ResourceAggregate),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(ResourceAggregate)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().setQuota(ROLE, ResourceAggregate())
 
   def test_forceTaskState(self):
     self.mock_thrift_client.forceTaskState(
         'taskid',
-        IgnoreArg(),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().forceTaskState('taskid', ScheduleStatus.LOST)
 
   def test_performBackup(self):
-    self.mock_thrift_client.performBackup(SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.performBackup().AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().performBackup()
 
   def test_listBackups(self):
-    self.mock_thrift_client.listBackups(SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.listBackups().AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().listBackups()
 
   def test_stageRecovery(self):
     self.mock_thrift_client.stageRecovery(
-        IsA(TaskQuery),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().stageRecovery(TaskQuery())
 
   def test_queryRecovery(self):
     self.mock_thrift_client.queryRecovery(
-      IsA(TaskQuery),
-      SESSION).AndReturn(DEFAULT_RESPONSE)
+      IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().queryRecovery(TaskQuery())
 
   def test_deleteRecoveryTasks(self):
     self.mock_thrift_client.deleteRecoveryTasks(
-        IsA(TaskQuery),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().deleteRecoveryTasks(TaskQuery())
 
   def test_commitRecovery(self):
-    self.mock_thrift_client.commitRecovery(SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.commitRecovery().AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().commitRecovery()
 
   def test_unloadRecovery(self):
-    self.mock_thrift_client.unloadRecovery(SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.unloadRecovery().AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().unloadRecovery()
 
   def test_snapshot(self):
-    self.mock_thrift_client.snapshot(SESSION).AndReturn(DEFAULT_RESPONSE)
+    self.mock_thrift_client.snapshot().AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().snapshot()
 
   def test_rewriteConfigs(self):
     self.mock_thrift_client.rewriteConfigs(
-        IsA(RewriteConfigsRequest),
-        SESSION).AndReturn(DEFAULT_RESPONSE)
+        IsA(RewriteConfigsRequest)).AndReturn(DEFAULT_RESPONSE)
     self.mox.ReplayAll()
     self.make_scheduler_proxy().rewriteConfigs(RewriteConfigsRequest())
 
@@ -438,7 +410,7 @@ class TestSchedulerClient(unittest.TestCase):
     mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client
     client.get.return_value = mock_scheduler_client
 
-    proxy = TestSchedulerProxy(Cluster(name='local'))
+    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
     proxy.killTasks(TaskQuery(), None)
 
     assert mock_thrift_client.killTasks.call_count == 3
@@ -449,7 +421,7 @@ class TestSchedulerClient(unittest.TestCase):
     mock_scheduler_client = 
mock.create_autospec(spec=scheduler_client.SchedulerClient,
                                                  instance=True)
     client.get.return_value = mock_scheduler_client
-    proxy = TestSchedulerProxy(Cluster(name='local'))
+    proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
 
     # unknown, transient connection error
     mock_scheduler_client.get_thrift_client.side_effect = RuntimeError

http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/python/apache/aurora/client/api/test_updater.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_updater.py 
b/src/test/python/apache/aurora/client/api/test_updater.py
index b7b1653..5f54d22 100644
--- a/src/test/python/apache/aurora/client/api/test_updater.py
+++ b/src/test/python/apache/aurora/client/api/test_updater.py
@@ -53,7 +53,6 @@ from gen.apache.aurora.api.ttypes import (
     ScheduledTask,
     ScheduleStatusResult,
     ServerInfo,
-    SessionKey,
     TaskConfig,
     TaskConstraint,
     TaskQuery,
@@ -149,15 +148,12 @@ class UpdaterTest(TestCase):
     self._name = 'jimbob'
     self._env = 'test'
     self._job_key = JobKey(name=self._name, environment=self._env, 
role=self._role)
-    self._session_key = SessionKey()
     self._lock = 'test_lock'
     self._instance_watcher = MockObject(InstanceWatcher)
     self._job_monitor = MockObject(JobMonitor)
     self._scheduler_mux = FakeSchedulerMux()
     self._scheduler = MockObject(scheduler_client)
-    self._scheduler_proxy = FakeSchedulerProxy(Cluster(name='test-cluster'),
-                                               self._scheduler,
-                                               self._session_key)
+    self._scheduler_proxy = FakeSchedulerProxy(Cluster(name='test-cluster'), 
self._scheduler)
     self._quota_check = MockObject(QuotaCheck)
     self.init_updater(deepcopy(self.UPDATE_CONFIG))
     self._num_cpus = 1.0
@@ -215,7 +211,7 @@ class UpdaterTest(TestCase):
 
   def expect_cron_replace(self, job_config, response_code=ResponseCode.OK):
     resp = make_response(response_code)
-    self._scheduler.replaceCronTemplate(job_config, self._lock, 
self._session_key).AndReturn(resp)
+    self._scheduler.replaceCronTemplate(job_config, self._lock).AndReturn(resp)
 
   def expect_restart(self, instance_ids, response_code=None):
     for i in instance_ids:
@@ -224,8 +220,7 @@ class UpdaterTest(TestCase):
       self._scheduler.restartShards(
           self._job_key,
           [i],
-          self._lock,
-          self._session_key).AndReturn(response)
+          self._lock).AndReturn(response)
 
   def expect_kill(self,
       instance_ids,
@@ -238,8 +233,7 @@ class UpdaterTest(TestCase):
                         instanceIds=frozenset([int(i)]))
       self._scheduler.killTasks(
           query,
-          self._lock,
-          
self._session_key).InAnyOrder().AndReturn(make_response(response_code))
+          self._lock).InAnyOrder().AndReturn(make_response(response_code))
 
     self.expect_job_monitor(response_code, instance_ids, monitor_result, 
skip_monitor)
 
@@ -260,8 +254,7 @@ class UpdaterTest(TestCase):
           instanceIds=frozenset([int(i)]))
       self._scheduler.addInstances(
           add_config,
-          self._lock,
-          
self._session_key).InAnyOrder().AndReturn(make_response(response_code))
+          self._lock).InAnyOrder().AndReturn(make_response(response_code))
 
   def expect_update_instances(self, instance_ids, task_config):
     for i in instance_ids:
@@ -281,13 +274,12 @@ class UpdaterTest(TestCase):
   def expect_start(self, response_code=ResponseCode.OK):
     response = make_response(response_code)
     response.result = 
Result(acquireLockResult=AcquireLockResult(lock=self._lock))
-    self._scheduler.acquireLock(LockKey(job=self._job_key), 
self._session_key).AndReturn(response)
+    self._scheduler.acquireLock(LockKey(job=self._job_key)).AndReturn(response)
 
   def expect_finish(self, response_code=ResponseCode.OK):
     self._scheduler.releaseLock(
         self._lock,
-        LockValidation.CHECKED,
-        self._session_key).AndReturn(make_response(response_code))
+        LockValidation.CHECKED).AndReturn(make_response(response_code))
 
   def expect_quota_check(self,
       num_released,

http://git-wip-us.apache.org/repos/asf/aurora/blob/7b1c2b2f/src/test/python/apache/aurora/client/fake_scheduler_proxy.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/fake_scheduler_proxy.py 
b/src/test/python/apache/aurora/client/fake_scheduler_proxy.py
index 183d2d6..7a8f766 100644
--- a/src/test/python/apache/aurora/client/fake_scheduler_proxy.py
+++ b/src/test/python/apache/aurora/client/fake_scheduler_proxy.py
@@ -18,18 +18,14 @@ from apache.aurora.client.api.scheduler_client import 
SchedulerProxy
 
 
 class FakeSchedulerProxy(SchedulerProxy):
-  def __init__(self, cluster, scheduler, session_key):
+  def __init__(self, cluster, scheduler):
     SchedulerProxy.__init__(self, cluster, verbose=True)
     self._scheduler = scheduler
-    self._session_key = session_key
     self._lock = RLock()
     self._terminating = Event()
 
   def client(self):
     return self._scheduler
 
-  def session_key(self):
-    return self._session_key
-
   def terminate(self):
     return self._terminating.set()

Reply via email to