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())