class TestSchedulerProxyInjection(unittest.TestCase):
def setUp(self):
self.mox = Mox()
self.mox.StubOutClassWithMocks(AuroraAdmin, 'Client')
self.mox.StubOutClassWithMocks(scheduler_client, 'SchedulerClient')
self.mock_scheduler_client = self.mox.CreateMock(scheduler_client.SchedulerClient)
self.mock_thrift_client = self.mox.CreateMock(AuroraAdmin.Client)
scheduler_client.SchedulerClient.get(IgnoreArg(), verbose=IgnoreArg()).AndReturn(
self.mock_scheduler_client)
self.mock_scheduler_client.get_thrift_client().AndReturn(self.mock_thrift_client)
def tearDown(self):
self.mox.UnsetStubs()
self.mox.VerifyAll()
def make_scheduler_proxy(self):
return scheduler_client.SchedulerProxy(Cluster(name='local'))
def test_startCronJob(self):
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)).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)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().replaceCronTemplate(JobConfiguration(), Lock())
def test_scheduleCronJob(self):
self.mock_thrift_client.scheduleCronJob(
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)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().descheduleCronJob(JOB_KEY)
def test_populateJobConfig(self):
self.mock_thrift_client.populateJobConfig(IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().populateJobConfig(JobConfiguration())
def test_restartShards(self):
self.mock_thrift_client.restartShards(
IsA(JobKey),
IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().restartShards(JOB_KEY, set([0]))
def test_getTasksStatus(self):
self.mock_thrift_client.getTasksStatus(IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getTasksStatus(TaskQuery())
def test_getJobs(self):
self.mock_thrift_client.getJobs(IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobs(ROLE)
def test_killTasks(self):
self.mock_thrift_client.killTasks(
IgnoreArg(),
IgnoreArg(),
IsA(JobKey),
IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().killTasks(None, None, JobKey(), set([0]))
def test_getQuota(self):
self.mock_thrift_client.getQuota(IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getQuota(ROLE)
def test_addInstances(self):
self.mock_thrift_client.addInstances(
IsA(JobKey),
IgnoreArg(),
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)).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)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().releaseLock(Lock(), LockValidation())
def test_getJobUpdateSummaries(self):
self.mock_thrift_client.getJobUpdateSummaries(IsA(JobUpdateQuery)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobUpdateSummaries(JobUpdateQuery())
def test_getJobUpdateDetails(self):
self.mock_thrift_client.getJobUpdateDetails('update_id').AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobUpdateDetails('update_id')
def test_startJobUpdate(self):
self.mock_thrift_client.startJobUpdate(
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').AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().pauseJobUpdate('update_id')
def test_resumeJobUpdate(self):
self.mock_thrift_client.resumeJobUpdate(
'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').AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().abortJobUpdate('update_id')
def test_pulseJobUpdate(self):
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).AndRaise(
TRequestsTransport.AuthError())
self.mock_scheduler_client.get_failed_auth_message().AndReturn('failed auth')
self.mox.ReplayAll()
with pytest.raises(scheduler_client.SchedulerProxy.AuthError):
self.make_scheduler_proxy().killTasks(TaskQuery(), None, None)
class TestSchedulerProxyInjection(unittest.TestCase):
def setUp(self):
self.mox = Mox()
self.mox.StubOutClassWithMocks(AuroraAdmin, 'Client')
self.mox.StubOutClassWithMocks(scheduler_client, 'SchedulerClient')
self.mock_scheduler_client = self.mox.CreateMock(
scheduler_client.SchedulerClient)
self.mock_thrift_client = self.mox.CreateMock(AuroraAdmin.Client)
scheduler_client.SchedulerClient.get(IgnoreArg(),
verbose=IgnoreArg()).AndReturn(
self.mock_scheduler_client)
self.mock_scheduler_client.get_thrift_client().AndReturn(
self.mock_thrift_client)
def tearDown(self):
self.mox.UnsetStubs()
self.mox.VerifyAll()
def make_scheduler_proxy(self):
return TestSchedulerProxy(Cluster(name='local'))
def test_startCronJob(self):
self.mock_thrift_client.startCronJob(
IsA(JobKey), IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().startCronJob(JOB_KEY)
def test_createJob(self):
self.mock_thrift_client.createJob(
IsA(JobConfiguration), IsA(SessionKey)).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),
IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().replaceCronTemplate(
JobConfiguration(), Lock())
def test_scheduleCronJob(self):
self.mock_thrift_client.scheduleCronJob(
IsA(JobConfiguration), IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().scheduleCronJob(JobConfiguration())
def test_descheduleCronJob(self):
self.mock_thrift_client.descheduleCronJob(
IsA(JobKey), IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().descheduleCronJob(JOB_KEY)
def test_populateJobConfig(self):
self.mock_thrift_client.populateJobConfig(
IsA(JobConfiguration)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().populateJobConfig(JobConfiguration())
def test_restartShards(self):
self.mock_thrift_client.restartShards(
IsA(JobKey), IgnoreArg(),
IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().restartShards(JOB_KEY, set([0]))
def test_getTasksStatus(self):
self.mock_thrift_client.getTasksStatus(
IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getTasksStatus(TaskQuery())
def test_getJobs(self):
self.mock_thrift_client.getJobs(
IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobs(ROLE)
def test_killTasks(self):
self.mock_thrift_client.killTasks(
IsA(TaskQuery), IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().killTasks(TaskQuery())
def test_getQuota(self):
self.mock_thrift_client.getQuota(
IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getQuota(ROLE)
def test_api_version_mismatch(self):
resp = Response(serverInfo=ServerInfo(
thriftAPIVersion=THRIFT_API_VERSION + 1))
self.mock_thrift_client.getQuota(IgnoreArg()).AndReturn(resp)
self.mox.ReplayAll()
with pytest.raises(
scheduler_client.SchedulerProxy.ThriftInternalError):
self.make_scheduler_proxy().getQuota(ROLE)
def test_addInstances(self):
self.mock_thrift_client.addInstances(
IsA(JobKey), IgnoreArg(), IsA(Lock),
IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().addInstances(JobKey(), {}, Lock())
def test_acquireLock(self):
self.mock_thrift_client.acquireLock(
IsA(Lock), IsA(SessionKey)).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),
IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().releaseLock(Lock(), LockValidation())
def test_getJobUpdateSummaries(self):
self.mock_thrift_client.getJobUpdateSummaries(
IsA(JobUpdateQuery)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobUpdateSummaries(JobUpdateQuery())
def test_getJobUpdateDetails(self):
self.mock_thrift_client.getJobUpdateDetails('update_id').AndReturn(
DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobUpdateDetails('update_id')
def test_startJobUpdate(self):
self.mock_thrift_client.startJobUpdate(
IsA(JobUpdateRequest), IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().startJobUpdate(JobUpdateRequest())
def test_pauseJobUpdate(self):
self.mock_thrift_client.pauseJobUpdate(
'update_id', IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().pauseJobUpdate('update_id')
def test_resumeJobUpdate(self):
self.mock_thrift_client.resumeJobUpdate(
'update_id', IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().resumeJobUpdate('update_id')
def test_abortJobUpdate(self):
self.mock_thrift_client.abortJobUpdate(
'update_id', IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().abortJobUpdate('update_id')
def test_pulseJobUpdate(self):
self.mock_thrift_client.pulseJobUpdate(
'update_id', IsA(SessionKey)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().pulseJobUpdate('update_id')
class TestSchedulerProxyInjection(unittest.TestCase):
def setUp(self):
self.mox = Mox()
self.mox.StubOutClassWithMocks(AuroraAdmin, 'Client')
self.mox.StubOutClassWithMocks(scheduler_client, 'SchedulerClient')
self.mock_scheduler_client = self.mox.CreateMock(
scheduler_client.SchedulerClient)
self.mock_thrift_client = self.mox.CreateMock(AuroraAdmin.Client)
scheduler_client.SchedulerClient.get(IgnoreArg(),
verbose=IgnoreArg()).AndReturn(
self.mock_scheduler_client)
self.mock_scheduler_client.get_thrift_client().AndReturn(
self.mock_thrift_client)
version_resp = Response(responseCode=ResponseCode.OK)
version_resp.result = Result(getVersionResult=CURRENT_API_VERSION)
self.mock_thrift_client.getVersion().AndReturn(version_resp)
def tearDown(self):
self.mox.UnsetStubs()
self.mox.VerifyAll()
def make_scheduler_proxy(self):
return TestSchedulerProxy('local')
def test_startCronJob(self):
self.mock_thrift_client.startCronJob(IsA(JobKey), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().startCronJob(JOB_KEY)
def test_createJob(self):
self.mock_thrift_client.createJob(IsA(JobConfiguration),
IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().createJob(JobConfiguration())
def test_replaceCronTemplate(self):
self.mock_thrift_client.replaceCronTemplate(IsA(JobConfiguration),
IsA(Lock), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().replaceCronTemplate(
JobConfiguration(), Lock())
def test_populateJobConfig(self):
self.mock_thrift_client.populateJobConfig(IsA(JobConfiguration))
self.mox.ReplayAll()
self.make_scheduler_proxy().populateJobConfig(JobConfiguration())
def test_restartShards(self):
self.mock_thrift_client.restartShards(IsA(JobKey), IgnoreArg(),
IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().restartShards(JOB_KEY, set([0]))
def test_getTasksStatus(self):
self.mock_thrift_client.getTasksStatus(IsA(TaskQuery))
self.mox.ReplayAll()
self.make_scheduler_proxy().getTasksStatus(TaskQuery())
def test_getJobs(self):
self.mock_thrift_client.getJobs(IgnoreArg())
self.mox.ReplayAll()
self.make_scheduler_proxy().getJobs(ROLE)
def test_killTasks(self):
self.mock_thrift_client.killTasks(IsA(TaskQuery), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().killTasks(TaskQuery())
def test_getQuota(self):
self.mock_thrift_client.getQuota(IgnoreArg())
self.mox.ReplayAll()
self.make_scheduler_proxy().getQuota(ROLE)
def test_startMaintenance(self):
self.mock_thrift_client.startMaintenance(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().startMaintenance(Hosts())
def test_drainHosts(self):
self.mock_thrift_client.drainHosts(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().drainHosts(Hosts())
def test_maintenanceStatus(self):
self.mock_thrift_client.maintenanceStatus(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().maintenanceStatus(Hosts())
def test_endMaintenance(self):
self.mock_thrift_client.endMaintenance(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().endMaintenance(Hosts())
def test_getVersion(self):
self.mock_thrift_client.getVersion()
self.mox.ReplayAll()
self.make_scheduler_proxy().getVersion()
def test_addInstances(self):
self.mock_thrift_client.addInstances(IsA(JobKey), IgnoreArg(),
IsA(Lock), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().addInstances(JobKey(), {}, Lock())
def test_acquireLock(self):
self.mock_thrift_client.acquireLock(IsA(Lock), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().acquireLock(Lock())
def test_releaseLock(self):
self.mock_thrift_client.releaseLock(IsA(Lock), IsA(LockValidation),
IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().releaseLock(Lock(), LockValidation())